Force fields to have the same eltype

[gdalle]

In the following struct for example

struct DeviceSparseMatrixCSR{
    Tv,
    Ti<:Integer,
    RowPtrT<:AbstractVector,
    ColValT<:AbstractVector,
    NzValT<:AbstractVector,
} <: AbstractDeviceSparseMatrix{Tv,Ti}

is there a reason not to have

RowPtrT<:AbstractVector{Ti}
ColValT<:AbstractVector{Ti}
NzValT<:AbstractVector{Tv}

[albertomercurio]

I remember I had some problems when using Reactant.jl, because Tv could also be a TracedRNumber or something similar.

It definitely needs to be investigated, but for sure one should make RowPtrT and ColValT to be of the same type. What I'm not sure is if NzVaT can always have Tv eltype.

[gdalle]

But then I'm confused, what is the meaning of Tv if the nonzero values have a different type?

[albertomercurio]

The form of Reactant is TracedRNumber{Tv} still. So it is linked to it, just wrapped from Reactant. That' why I'm using _get_eltype like here:

DeviceSparseArrays.jl/src/matrix_coo.jl

Lines 71 to 73 in 98dfe9e

	Ti2 = _get_eltype(rowind)
	Tv2 = _get_eltype(nzval)
	DeviceSparseMatrixCOO{Tv2,Ti2,RowIndT,ColIndT,NzValT}(m, n, rowind, colind, nzval)

But again, I did it in the beginning as a quick fix. Maybe there is no need to do it.

[gdalle]

I think it would make more sense to have Tv be the actual element type, even if it is a Reactant number type. Are there cases where you truly needed to access the underlying floating point type for example?

[albertomercurio]

Ok, I'm giving a try again. I made PR #28, but as I said, this has some problems with Reactant.

using DeviceSparseArrays, Adapt, LinearAlgebra, SparseArrays
using CUDA
using Reactant

A_sparse = sprand(Float64, 100, 100, 0.1)
A_device = DeviceSparseMatrixCSC(A_sparse)
A_reactant = adapt(Reactant.ConcreteRArray, A_device)

tr_compiled = @compile tr(A_reactant)

ERROR: TypeError: in DeviceSparseMatrixCSC, in ColPtrT, expected ColPtrT<:AbstractVector{Int64}, got Type{Reactant.TracedRArray{Int64, 1}}
Stacktrace:
  [1] traced_type_inner(T::Type, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:750
  [2] traced_type(T::Type, ::Val{Reactant.ConcreteToTraced}, track_numbers::Type, sharding::Reactant.Sharding.NoSharding, runtime::Val{:PJRT})
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:886
  [3] make_tracer_unknown(seen::Reactant.OrderedIdDict{…}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{…})
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1048
  [4] make_tracer_unknown
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1025 [inlined]
  [5] #make_tracer#137
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1162 [inlined]
  [6] make_tracer
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1152 [inlined]
  [7] prepare_mlir_fn_args(args::Tuple{…}, name::String, concretein::Bool, toscalar::Bool, argprefix::Symbol, runtime::Val{…}, optimize_then_pad::Bool, do_transpose::Bool, input_shardings::Nothing, verify_arg_names::Nothing)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:450
  [8] make_mlir_fn(f::typeof(tr), args::Tuple{…}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:321
  [9] make_mlir_fn
    @ ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:275 [inlined]
 [10] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(tr), args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}, sdygroupidcache::Tuple{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, client::Reactant.XLA.PJRT.Client, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:1605
 [11] compile_mlir!
    @ ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:1567 [inlined]
 [12] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3513
 [13] compile_xla
    @ ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3485 [inlined]
 [14] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3589
 [15] top-level scope

Which works on main instead.

Somehow it is not promoting Tv and Ti.

[gdalle]

I think we might want to relax the constraint Ti <: Integer, because of the following subtype relations:

https://github.com/EnzymeAD/Reactant.jl/blob/f6643eb3fba75d1fbde8f5adca9c3ac88f841cf9/src/Types.jl#L67

https://github.com/EnzymeAD/Reactant.jl/blob/f6643eb3fba75d1fbde8f5adca9c3ac88f841cf9/src/Types.jl#L41

https://github.com/EnzymeAD/Reactant.jl/blob/f6643eb3fba75d1fbde8f5adca9c3ac88f841cf9/src/Types.jl#L1C15-L1C44

In other words, TracedRArray{Int64, 1} is not a vector of integers but a vector of numbers.

[gdalle]

Although this example tells me that we probably want to be careful with Reactant too (I'm not an expert at all in the tracing mechanism):

JuliaDiff/DifferentiationInterface.jl#918 (comment)

[albertomercurio]

Nice catch, but it still fails. I found also the following code to fail. I will open an issue on Reactant to ask help.

julia> struct MyStruct{T, VT<:AbstractVector{T}}
           data::VT
       end

julia> test(x::MyStruct) = sum(x.data)
test (generic function with 1 method)

julia> x = MyStruct(rand(10))
MyStruct{Float64, Vector{Float64}}([0.9356372116439775, 0.8146232142268686, 0.8646666386255512, 0.15223146438404445, 0.5357206445516378, 0.12556851518402812, 0.014343067486348171, 0.6245919019686935, 0.9909782907741145, 0.3270593445135602])

julia> test(x)
5.385420293358824

julia> dx = MyStruct(Reactant.ConcreteRArray(x.data))
MyStruct{Float64, ConcretePJRTArray{Float64, 1, 1}}(ConcretePJRTArray{Float64, 1, 1}([0.9356372116439775, 0.8146232142268686, 0.8646666386255512, 0.15223146438404445, 0.5357206445516378, 0.12556851518402812, 0.014343067486348171, 0.6245919019686935, 0.9909782907741145, 0.3270593445135602]))

julia> test_compiled = @compile test(dx)
ERROR: TypeError: in MyStruct, in VT, expected VT<:AbstractVector{Float64}, got Type{Reactant.TracedRArray{Float64, 1}}
Stacktrace:
  [1] traced_type_inner(T::Type, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:750
  [2] traced_type(T::Type, ::Val{Reactant.ConcreteToTraced}, track_numbers::Type, sharding::Reactant.Sharding.NoSharding, runtime::Val{:PJRT})
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:886
  [3] make_tracer_unknown(seen::Reactant.OrderedIdDict{…}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{…})
    @ Reactant ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1048
  [4] make_tracer_unknown
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1025 [inlined]
  [5] #make_tracer#137
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1162 [inlined]
  [6] make_tracer
    @ ~/.julia/packages/Reactant/8rzTQ/src/Tracing.jl:1152 [inlined]
  [7] prepare_mlir_fn_args(args::Tuple{…}, name::String, concretein::Bool, toscalar::Bool, argprefix::Symbol, runtime::Val{…}, optimize_then_pad::Bool, do_transpose::Bool, input_shardings::Nothing, verify_arg_names::Nothing)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:450
  [8] make_mlir_fn(f::typeof(test), args::Tuple{…}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:321
  [9] make_mlir_fn
    @ ~/.julia/packages/Reactant/8rzTQ/src/TracedUtils.jl:275 [inlined]
 [10] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(test), args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}, sdygroupidcache::Tuple{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, client::Reactant.XLA.PJRT.Client, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:1605
 [11] compile_mlir!
    @ ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:1567 [inlined]
 [12] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3513
 [13] compile_xla
    @ ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3485 [inlined]
 [14] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:3589
 [15] top-level scope
    @ ~/.julia/packages/Reactant/8rzTQ/src/Compiler.jl:2658
Some type information was truncated. Use `show(err)` to see complete types.

[gdalle]

Yeah I'm curious to see what they say. If we don't understand that, we won't understand #28