Merge pull request #103 from EnzymeAD/ap/concat

feat: handle cat/hcat/vcat

Author: avik-pal

src/TracedRArray.jl

@@ -758,3 +758,23 @@ function _copyto!(dest::TracedRArray, bc::Broadcasted)
     dest.mlir_data = res.mlir_data
     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."
+    catdims = Base.dims2cat(dims)
+    shape = Base.cat_size_shape(catdims, A, Bs...)
+    RT = Base.promote_eltype(A, Bs...)
+    Res = TracedRArray{RT,length(shape)}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.concatenate(
+                [A.mlir_data, [B.mlir_data for B in Bs]...];
+                result_0=MLIR.IR.TensorType(shape, MLIR.IR.Type(RT)),
+                dimension=D - 1, # stablehlo expects this to be zero-indexed
+            ),
+            1,
+        ),
+        shape,
+    )
+    return Res
+end

test/basic.jl

@@ -195,3 +195,20 @@ end
     @test var_fn3(x) ≈ var_fn3_compiled(x_ca)
     @test var_fn4(x) ≈ var_fn4_compiled(x_ca)
 end
+
+@testset "concatenation" begin
+    x = ones(2, 4, 3)
+    x_concrete = Reactant.to_rarray(x)
+
+    cat1(x) = vcat(x, x, x)
+    cat2(x) = hcat(x, x, x)
+    cat3(x) = cat(x, x, x; dims=Val(3))
+
+    cat1_compiled = Reactant.compile(cat1, (x_concrete,))
+    cat2_compiled = Reactant.compile(cat2, (x_concrete,))
+    cat3_compiled = Reactant.compile(cat3, (x_concrete,))
+
+    @test cat1(x) ≈ cat1_compiled(x_concrete)
+    @test cat2(x) ≈ cat2_compiled(x_concrete)
+    @test cat3(x) ≈ cat3_compiled(x_concrete)
+end