Generalize Base._cat to non-Val, typed Base._cat_t and implement typed_hcat, typed_vcat, typed_hvcat, typed_hvncat

src/TracedRArray.jl

@@ -761,32 +761,36 @@ function _copyto!(dest::TracedRArray, bc::Broadcasted)
     return dest
 end
 
-function Base._cat(dims::Val{D}, A::TracedRArray{T,N}, Bs::TracedRArray...) where {T,N,D}
-    @assert D isa Integer "Support for non-integer dimensions is not implemented yet."
+dispatch_val(x) = x
+dispatch_val(::Val{D}) where {D} = D
 
-    # MLIR expects the dimension `D` to be ≤ the rank of the input tensors
-    A = maybe_expand_dims(A, dims)
-    Bs = maybe_expand_dims.(Bs, (dims,))
+function Base._cat_t(dims, ::Type{T}, X::TracedRArray...) where {T}
+    dims = dispatch_val(dims)
+    @assert dims isa Integer "Support for non-integer dimensions is not implemented yet."
+
+    # MLIR expects the dimension `dims` to be ≤ the rank of the input tensors
+    X = maybe_expand_dims.(X, (dims,))
 
     catdims = Base.dims2cat(dims)
-    shape = Base.cat_size_shape(catdims, A, Bs...)
-    RT = Base.promote_eltype(A, Bs...)
-    Res = TracedRArray{RT,length(shape)}(
+    shape = Base.cat_size_shape(catdims, X...)
+    RT = Base.promote_eltype(T, X...)
+    return TracedRArray{RT,length(shape)}(
         (),
         MLIR.IR.result(
+            # TODO maybe we should do some conversion?
             MLIR.Dialects.stablehlo.concatenate(
-                [A.mlir_data, [B.mlir_data for B in Bs]...];
+                collect(get_mlir_data.(X));
                 result_0=MLIR.IR.TensorType(shape, MLIR.IR.Type(RT)),
-                dimension=D - 1, # stablehlo expects this to be zero-indexed
+                dimension=dims - 1, # stablehlo expects this to be zero-indexed
             ),
             1,
         ),
         shape,
     )
-    return Res
 end
 
-function maybe_expand_dims(x::AbstractArray{T,N}, ::Val{D}) where {T,N,D}
-    D ≤ N && return x
-    return reshape(x, ntuple(i -> i ≤ N ? size(x, i) : 1, Val(D)))
+function maybe_expand_dims(x::AbstractArray{T,N}, dims) where {T,N}
+    dims = dispatch_val(dims)
+    dims ≤ N && return x
+    return reshape(x, ntuple(i -> i ≤ N ? size(x, i) : 1, dims))
 end

src/Tracing.jl

@@ -380,7 +380,7 @@ function make_tracer(
     if haskey(seen, prev)
         return seen[prev]
     end
-    if mode == ArrayToConcrete && eltype(RT) <: AbstractFloat
+    if mode == ArrayToConcrete && eltype(RT) <: Union{AbstractFloat, Integer}
         return seen[prev] = ConcreteRArray(prev)
     end
     TT = traced_type(eltype(RT), (), Val(mode))

test/basic.jl

@@ -210,6 +210,112 @@ end
 end
 
 @testset "concatenation" begin
+    @testset "0-dim" begin
+        x = fill(true)
+        x_concrete = Reactant.to_rarray(x)
+
+        # NOTE [,,,] is a call to `vect`, not `*cat`
+        # f = Reactant.compile((x_concrete,)) do x
+        #     return [x, x, x]
+        # end
+        # @test f(x_concrete) ≈ ones(3)
+
+        # vcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x; x; x]
+        end
+        @test f(x_concrete) == ones(Bool, 3)
+        @test eltype(f(x_concrete)) === Bool
+
+        # hcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x x x]
+        end
+        @test f(x_concrete) == ones(Bool, 1, 3)
+        @test eltype(f(x_concrete)) === Bool
+
+        # hvcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x x x; x x x]
+        end
+        @test f(x_concrete) == ones(Bool, 2, 3)
+        @test eltype(f(x_concrete)) === Bool
+
+        # typed_vcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x; x; x]
+        end
+        @test f(x_concrete) == ones(Int, 3)
+        @test eltype(f(x_concrete)) === Int
+
+        # typed_hcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x; x; x]
+        end
+        @test f(x_concrete) == ones(Int, 1, 3)
+        @test eltype(f(x_concrete)) === Int
+
+        # typed_hvcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x x x; x x x]
+        end
+        @test f(x_concrete) == ones(Int, 2, 3)
+        @test eltype(f(x_concrete)) === Int
+    end
+
+    @testset "1-dim" begin
+        x = ones(Bool, 2)
+        x_concrete = Reactant.to_rarray(x)
+
+        # NOTE [,,,] is a call to `vect`, not `*cat`
+        # f = Reactant.compile((x_concrete,)) do x
+        #     return [x, x, x]
+        # end
+        # @test f(x_concrete) ≈ ones(3)
+
+        # vcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x; x; x]
+        end
+        @test f(x_concrete) == ones(Bool, 6)
+        @test eltype(f(x_concrete)) === Bool
+
+        # hcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x x x]
+        end
+        @test f(x_concrete) == ones(Bool, 2, 3)
+        @test eltype(f(x_concrete)) === Bool
+
+        # hvcat
+        f = Reactant.compile((x_concrete,)) do x
+            return [x x x; x x x]
+        end
+        @test f(x_concrete) == ones(Bool, 4, 3)
+        @test eltype(f(x_concrete)) === Bool
+
+        # typed_vcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x; x; x]
+        end
+        @test f(x_concrete) == ones(Int, 6)
+        @test eltype(f(x_concrete)) === Int
+
+        # typed_hcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x; x; x]
+        end
+        @test f(x_concrete) == ones(Int, 2, 3)
+        @test eltype(f(x_concrete)) === Int
+
+        # typed_hvcat
+        f = Reactant.compile((x_concrete,)) do x
+            return Int[x x x; x x x]
+        end
+        @test f(x_concrete) == ones(Int, 4, 3)
+        @test eltype(f(x_concrete)) === Int
+    end
+
     x = ones(2, 4, 3)
     x_concrete = Reactant.to_rarray(x)