Implement conj, conj! for TracedRArray (#169)

* Implement `conj` for `TracedRArray`

* Generalize `adjoint` for `TracedRVecOrMat`

* Implement `conj` for `TracedRNumber`

* Implement `conj!` for `TracedRArray`

* Implement `DenseElementsAttribute` for array of `Complex`

* Implement `Base.real`, `Base.imag` for `TracedRNumber`

* Fix typo

* Implement `Base.real`, `Base.imag` for `TracedRArray`

* Fix pointer length in `MLIR.IR.DenseElementsAttribute` on `Complex{T}`

* Move complex tests to new file

* Fix `ConcreteRArray` constructor on `Number`

* Fix `to_rarray` on primitive number types

* Fix `conj` tests on numbers

* Write tests for `real`, `imag`

* Remove duplicated method

* Update src/TracedRNumber.jl

Co-authored-by: Paul Berg <[email protected]>

* fix `image` on `TracedRArray` of reals

---------

Co-authored-by: Paul Berg <[email protected]>

Author: mofeing

src/ConcreteRArray.jl

@@ -8,7 +8,7 @@ mutable struct ConcreteRArray{T,N} <: RArray{T,N}
     shape::NTuple{N,Int}
 end
 
-ConcreteRArray(data::T) where {T<:Number} = ConcreteRArray{T,0}(data, ())
+ConcreteRArray(data::T) where {T<:Number} = ConcreteRArray(fill(data))
 
 Adapt.adapt_storage(::Type{T}, x::AbstractArray) where {T<:ConcreteRArray} = T(x)
 

src/TracedRArray.jl

@@ -188,11 +188,62 @@ function Base.permutedims(A::AnyTracedRArray{T,N}, perm) where {T,N}
     )
 end
 
+Base.conj(A::TracedRArray) = A
+function Base.conj(A::TracedRArray{T,N}) where {T<:Complex,N}
+    return TracedRArray{T,N}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.chlo.conj(
+                A.mlir_data; result=mlir_type(TracedRArray{T,N}, size(A))
+            ),
+            1,
+        ),
+        size(A),
+    )
+end
+
+Base.conj!(A::TracedRArray) = A
+function Base.conj!(A::TracedRArray{T,N}) where {T<:Complex,N}
+    A.mlir_data = MLIR.IR.result(
+        MLIR.Dialects.chlo.conj(A.mlir_data; result=mlir_type(TracedRArray{T,N}, size(A))),
+        1,
+    )
+    return A
+end
+
+Base.real(A::TracedRArray) = A
+function Base.real(A::TracedRArray{Complex{T},N}) where {T,N}
+    return TracedRArray{T,N}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.real(
+                A.mlir_data; result=mlir_type(TracedRArray{T,N}, size(A))
+            ),
+            1,
+        ),
+        size(A),
+    )
+end
+
+Base.imag(A::TracedRArray) = zero(A)
+function Base.imag(A::TracedRArray{Complex{T},N}) where {T,N}
+    return TracedRArray{T,N}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.imag(
+                A.mlir_data; result=mlir_type(TracedRArray{T,N}, size(A))
+            ),
+            1,
+        ),
+        size(A),
+    )
+end
+
 function Base.transpose(A::AnyTracedRVecOrMat)
     A = ndims(A) == 1 ? reshape(A, :, 1) : A
     return permutedims(A, (2, 1))
 end
-Base.adjoint(A::AnyTracedRVecOrMat{<:Real}) = transpose(A)
+Base.adjoint(A::AnyTracedRVecOrMat) = conj(transpose(A))
 
 function promote_to(::Type{TracedRArray{T,N}}, rhs) where {T,N}
     if isa(rhs, TracedRArray)

src/TracedRNumber.jl

@@ -199,6 +199,36 @@ for (jlop, hloop) in (
     end
 end
 
+Base.conj(x::TracedRNumber) = x
+function Base.conj(x::TracedRNumber{T}) where {T<:Complex}
+    return TracedRNumber{T}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.chlo.conj(x.mlir_data; result=mlir_type(TracedRNumber{T})), 1
+        ),
+    )
+end
+
+Base.real(x::TracedRNumber) = x
+function Base.real(x::TracedRNumber{Complex{T}}) where {T}
+    return TracedRNumber{T}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.real(x.mlir_data; result=mlir_type(TracedRNumber{T})), 1
+        ),
+    )
+end
+
+Base.imag(x::TracedRNumber) = zero(x)
+function Base.imag(x::TracedRNumber{Complex{T}}) where {T}
+    return TracedRNumber{T}(
+        (),
+        MLIR.IR.result(
+            MLIR.Dialects.stablehlo.imag(x.mlir_data; result=mlir_type(TracedRNumber{T})), 1
+        ),
+    )
+end
+
 # XXX: Enzyme-MLIR doesn't have `abs` adjoint defined
 Base.abs2(x::TracedRNumber{<:Real}) = x^2
 

src/Tracing.jl

@@ -492,3 +492,5 @@ end
 @inline function to_rarray(@nospecialize(x))
     return make_tracer(OrderedIdDict(), x, (), Reactant.ArrayToConcrete)
 end
+
+to_rarray(x::ReactantPrimitive) = ConcreteRArray(x)

src/mlir/IR/Attribute.jl

@@ -567,8 +567,10 @@ end
 function DenseElementsAttribute(values::AbstractArray{<:Complex})
     shaped_type = TensorType(size(values), Type(eltype(values)))
     # TODO: row major
-    Attribute(
-        API.mlirDenseElementsAttrRawBufferGet(shaped_type, length(values) * sizeof(eltype(values)), values)
+    return Attribute(
+        API.mlirDenseElementsAttrRawBufferGet(
+            shaped_type, length(values) * Base.elsize(values), values
+        ),
     )
 end
 

src/mlir/IR/Type.jl

@@ -85,13 +85,6 @@ Creates a signless integer type of the given bitwidth in the context. The type i
 Type(T::Core.Type{<:Integer}; context::Context=context()) =
     Type(API.mlirIntegerTypeGet(context, sizeof(T) * 8))
 
-"""
-    Type(T::Core.Type{<:Complex}; context=context())
-
-Creates a complex type with the given element type.
-"""
-Type(T::Core.Type{<:Complex}; context=context()) = Type(API.mlirComplexTypeGet(Type(T(im) |> real |> typeof)))
-
 """
     Type(T::Core.Type{<:Signed}; context=context()
 

test/complex.jl

@@ -0,0 +1,89 @@
+using Test
+using Reactant
+
+@testset "conj" begin
+    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile conj(x_concrete)
+        @test only(f(x_concrete)) == conj(x)
+    end
+
+    @testset "$(typeof(x))" for x in (
+        fill(1.0 + 2.0im),
+        fill(1.0),
+        [1.0 + 2.0im; 3.0 + 4.0im],
+        [1.0; 3.0],
+        [1.0 + 2.0im 3.0 + 4.0im],
+        [1.0 2.0],
+        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
+        [1.0 3.0; 5.0 7.0],
+    )
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile conj(x_concrete)
+        @test f(x_concrete) == conj(x)
+    end
+end
+
+@testset "conj!" begin
+    @testset "$(typeof(x))" for x in (
+        fill(1.0 + 2.0im),
+        fill(1.0),
+        [1.0 + 2.0im; 3.0 + 4.0im],
+        [1.0; 3.0],
+        [1.0 + 2.0im 3.0 + 4.0im],
+        [1.0 2.0],
+        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
+        [1.0 3.0; 5.0 7.0],
+    )
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile conj!(x_concrete)
+        @test f(x_concrete) == conj(x)
+        @test x_concrete == conj(x)
+    end
+end
+
+@testset "real" begin
+    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile real(x_concrete)
+        @test only(f(x_concrete)) == real(x)
+    end
+
+    @testset "$(typeof(x))" for x in (
+        fill(1.0 + 2.0im),
+        fill(1.0),
+        [1.0 + 2.0im; 3.0 + 4.0im],
+        [1.0; 3.0],
+        [1.0 + 2.0im 3.0 + 4.0im],
+        [1.0 2.0],
+        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
+        [1.0 3.0; 5.0 7.0],
+    )
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile real(x_concrete)
+        @test f(x_concrete) == real(x)
+    end
+end
+
+@testset "imag" begin
+    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile imag(x_concrete)
+        @test only(f(x_concrete)) == imag(x)
+    end
+
+    @testset "$(typeof(x))" for x in (
+        fill(1.0 + 2.0im),
+        fill(1.0),
+        [1.0 + 2.0im; 3.0 + 4.0im],
+        [1.0; 3.0],
+        [1.0 + 2.0im 3.0 + 4.0im],
+        [1.0 2.0],
+        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
+        [1.0 3.0; 5.0 7.0],
+    )
+        x_concrete = Reactant.to_rarray(x)
+        f = @compile imag(x_concrete)
+        @test f(x_concrete) == imag(x)
+    end
+end

test/runtests.jl

@@ -46,6 +46,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Layout" include("layout.jl")
         @safetestset "Tracing" include("tracing.jl")
         @safetestset "Basic" include("basic.jl")
+        @safetestset "Complex" include("complex.jl")
         @safetestset "Broadcast" include("bcast.jl")
         @safetestset "Struct" include("struct.jl")
         @safetestset "Closure" include("closure.jl")