refactor: delegate simple opts to mlir/base julia (#1298)

Author: avik-pal

src/Compiler.jl

@@ -728,6 +728,7 @@ function optimization_passes(;
         "concat_reshape_reduce",
         "concat_elementwise",
         "reduce_reduce",
+        "conj_real",
         # TODO we want to enable but may cause an infinite compile time
         # "concat_to_onedim_dusslice",
     ]

src/Ops.jl

@@ -304,7 +304,7 @@ end
 
 @noinline function conj(
     x::TracedRArray{T,N}; location=mlir_stacktrace("conj", @__FILE__, @__LINE__)
-) where {T<:Complex,N}
+) where {T,N}
     res = MLIR.IR.result(
         chlo.conj(x.mlir_data; result=mlir_type(TracedRArray{T,N}, size(x)), location)
     )
@@ -313,7 +313,7 @@ end
 
 @noinline function conj(
     x::TracedRNumber{T}; location=mlir_stacktrace("conj", @__FILE__, @__LINE__)
-) where {T<:Complex}
+) where {T}
     res = MLIR.IR.result(
         chlo.conj(x.mlir_data; result=mlir_type(TracedRArray{T,0}, ()), location)
     )
@@ -597,39 +597,47 @@ end
 end
 
 @noinline function real(
-    x::TracedRArray{Complex{T},N}; location=mlir_stacktrace("real", @__FILE__, @__LINE__)
+    x::TracedRArray{T,N}; location=mlir_stacktrace("real", @__FILE__, @__LINE__)
 ) where {T,N}
     res = MLIR.IR.result(
-        stablehlo.real(x.mlir_data; result=mlir_type(TracedRArray{T,N}, size(x)), location)
+        stablehlo.real(
+            x.mlir_data; result=mlir_type(TracedRArray{Base.real(T),N}, size(x)), location
+        ),
     )
-    return TracedRArray{T,N}((), res, size(x))
+    return TracedRArray{Base.real(T),N}((), res, size(x))
 end
 
 @noinline function real(
-    x::TracedRNumber{Complex{T}}; location=mlir_stacktrace("real", @__FILE__, @__LINE__)
+    x::TracedRNumber{T}; location=mlir_stacktrace("real", @__FILE__, @__LINE__)
 ) where {T}
     res = MLIR.IR.result(
-        stablehlo.real(x.mlir_data; result=mlir_type(TracedRArray{T,0}, ()), location)
+        stablehlo.real(
+            x.mlir_data; result=mlir_type(TracedRArray{Base.real(T),0}, ()), location
+        ),
     )
-    return TracedRNumber{T}((), res)
+    return TracedRNumber{Base.real(T)}((), res)
 end
 
 @noinline function imag(
-    x::TracedRArray{Complex{T},N}; location=mlir_stacktrace("imag", @__FILE__, @__LINE__)
+    x::TracedRArray{T,N}; location=mlir_stacktrace("imag", @__FILE__, @__LINE__)
 ) where {T,N}
     res = MLIR.IR.result(
-        stablehlo.imag(x.mlir_data; result=mlir_type(TracedRArray{T,N}, size(x)), location)
+        stablehlo.imag(
+            x.mlir_data; result=mlir_type(TracedRArray{Base.real(T),N}, size(x)), location
+        ),
     )
-    return TracedRArray{T,N}((), res, size(x))
+    return TracedRArray{Base.real(T),N}((), res, size(x))
 end
 
 @noinline function imag(
-    x::TracedRNumber{Complex{T}}; location=mlir_stacktrace("imag", @__FILE__, @__LINE__)
+    x::TracedRNumber{T}; location=mlir_stacktrace("imag", @__FILE__, @__LINE__)
 ) where {T}
     res = MLIR.IR.result(
-        stablehlo.imag(x.mlir_data; result=mlir_type(TracedRArray{T,0}, ()), location)
+        stablehlo.imag(
+            x.mlir_data; result=mlir_type(TracedRArray{Base.real(T),0}, ()), location
+        ),
     )
-    return TracedRNumber{T}((), res)
+    return TracedRNumber{Base.real(T)}((), res)
 end
 
 function bitcast_convert(
@@ -679,7 +687,7 @@ end
         end
     elseif type == "IRFFT"
         @assert T <: Complex
-        Tout = Base.real(T)
+        Tout = Base.Base.real(T)
         rsize = let rsize = collect(Int64, size(x))
             rsize[(end - Base.length(length) + 1):end] = length
             Tuple(rsize)

src/TracedRArray.jl

@@ -511,22 +511,6 @@ function Base.permutedims(A::AnyTracedRArray{T,N}, perm) where {T,N}
     return Ops.transpose(materialize_traced_array(A), Int64[perm...])
 end
 
-Base.conj(A::AnyTracedRArray) = A
-Base.conj(A::AnyTracedRArray{<:Complex}) = Ops.conj(materialize_traced_array(A))
-
-Base.conj!(A::AnyTracedRArray) = A
-
-function Base.conj!(A::AnyTracedRArray{<:Complex})
-    TracedUtils.set_mlir_data!(A, Ops.conj(materialize_traced_array(A)).mlir_data)
-    return A
-end
-
-Base.real(A::AnyTracedRArray) = A
-Base.real(A::AnyTracedRArray{<:Complex}) = Ops.real(materialize_traced_array(A))
-
-Base.imag(A::AnyTracedRArray) = zero(A)
-Base.imag(A::AnyTracedRArray{<:Complex}) = Ops.imag(materialize_traced_array(A))
-
 TracedUtils.promote_to(::Type{TracedRArray{T,N}}, rhs) where {T,N} = TracedRArray{T,N}(rhs)
 function TracedUtils.promote_to(::TracedRArray{T,N}, rhs) where {T,N}
     return TracedUtils.promote_to(TracedRArray{T,N}, rhs)

src/TracedRNumber.jl

@@ -393,6 +393,9 @@ for (jlop, hloop) in (
     (:(Base.atan), :atan),
     (:(Base.atanh), :atanh),
     (:(Base.sign), :sign),
+    (:(Base.conj), :conj),
+    (:(Base.real), :real),
+    (:(Base.imag), :imag),
 )
     @eval $(jlop)(@nospecialize(lhs::TracedRNumber)) = Ops.$(hloop)(lhs)
 end
@@ -404,18 +407,9 @@ end
 
 Base.sincospi(x::TracedRNumber{T}) where {T} = Ops.sine(T(π) * x), Ops.cosine(T(π) * x)
 
-Base.conj(x::TracedRNumber) = x
-Base.conj(x::TracedRNumber{<:Complex}) = Ops.conj(x)
-
-Base.real(x::TracedRNumber) = x
-Base.real(x::TracedRNumber{<:Complex}) = Ops.real(x)
-
 Base.isreal(::TracedRNumber) = false
 Base.isreal(::TracedRNumber{<:Real}) = true
 
-Base.imag(x::TracedRNumber) = zero(x)
-Base.imag(x::TracedRNumber{<:Complex}) = Ops.imag(x)
-
 Base.iseven(x::TracedRNumber) = iseven(real(x))
 function Base.iseven(x::TracedRNumber{<:Real})
     return iszero(