feat: define a @opcall macro for better debug info (#1574)

* feat: define a @opcall macro for better debug info

* fix: missing location in Ops.select

* chore: run formatter

* fix: missing location in Ops.clamp

Author: avik-pal

ext/ReactantAbstractFFTsExt.jl

@@ -2,6 +2,7 @@ module ReactantAbstractFFTsExt
 
 using AbstractFFTs: AbstractFFTs
 using Reactant: Reactant, MLIR, Ops, AnyTracedRArray, TracedRArray, TracedUtils
+using Reactant.Ops: @opcall
 
 function __permutation_to_move_dims_to_end(dims, N::Integer)
     perm = [i for i in 1:N if i ∉ Set(dims)]
@@ -21,18 +22,20 @@ for op in (:rfft, :fft, :ifft)
 
         fft_lengths = Int64[size(x, dim) for dim in reverse(dims)]
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
-            return Ops.fft(
+            return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
                 type=$(uppercase(string(op))),
                 length=fft_lengths,
             )
         end
         perm = __permutation_to_move_dims_to_end(dims, ndims(x))
         return permutedims(
-            Ops.fft(
-                TracedUtils.materialize_traced_array(permutedims(x, perm));
-                type=$(uppercase(string(op))),
-                length=fft_lengths,
+            @opcall(
+                fft(
+                    TracedUtils.materialize_traced_array(permutedims(x, perm));
+                    type=$(uppercase(string(op))),
+                    length=fft_lengths,
+                )
             ),
             invperm(perm),
         )
@@ -48,7 +51,7 @@ for op in (:irfft,)
         fft_lengths = vcat(Int64[size(x, dim) for dim in reverse(dims[2:end])], d)
 
         if __is_valid_stablehlo_fft_dims(dims, ndims(x))
-            return Ops.fft(
+            return @opcall fft(
                 TracedUtils.materialize_traced_array(x);
                 type=$(uppercase(string(op))),
                 length=fft_lengths,
@@ -57,10 +60,12 @@ for op in (:irfft,)
 
         perm = __permutation_to_move_dims_to_end(dims, ndims(x))
         return permutedims(
-            Ops.fft(
-                TracedUtils.materialize_traced_array(permutedims(x, perm));
-                type=$(uppercase(string(op))),
-                length=fft_lengths,
+            @opcall(
+                fft(
+                    TracedUtils.materialize_traced_array(permutedims(x, perm));
+                    type=$(uppercase(string(op))),
+                    length=fft_lengths,
+                )
             ),
             invperm(perm),
         )

ext/ReactantCUDAExt.jl

@@ -10,6 +10,8 @@ import KernelAbstractions as KA
 using LLVM: LLVM
 using Libdl
 
+using Reactant.Ops: @opcall
+
 const ReactantKernelAbstractionsExt = Base.get_extension(
     Reactant, :ReactantKernelAbstractionsExt
 )
@@ -469,7 +471,7 @@ function Adapt.adapt_storage(ka::ReactantKernelAdaptor, xs::DenseCuArray)
     return Adapt.adapt_storage(ka, Array(xs))
 end
 function Adapt.adapt_storage(ka::ReactantKernelAdaptor, xs::Array)
-    return Adapt.adapt_storage(ka, Reactant.Ops.constant(xs))
+    return Adapt.adapt_storage(ka, @opcall(constant(xs)))
 end
 function Adapt.adapt_structure(
     to::ReactantKernelAdaptor, bc::Broadcast.Broadcasted{Style,<:Any,Type{T}}

ext/ReactantNNlibExt/Implementations.jl

@@ -3,7 +3,7 @@ for (jlop, hloop) in (
     (:(NNlib.sigmoid_fast), :logistic),
     (:(NNlib.sigmoid), :logistic),
 )
-    @eval $(jlop)(x::TracedRNumber) = Ops.$(hloop)(x)
+    @eval $(jlop)(x::TracedRNumber) = @opcall $(hloop)(x)
 end
 
 function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) where {T,N}
@@ -89,7 +89,7 @@ function overloaded_conv!(
 
     weight = W
     if !flipkernel
-        weight = Reactant.Ops.reverse(weight; dimensions=kernel_spatial_dims)
+        weight = @opcall reverse(weight; dimensions=kernel_spatial_dims)
     end
 
     conv = Reactant.MLIR.Dialects.stablehlo.convolution(
@@ -211,7 +211,7 @@ function overloaded_∇conv_filter!(
 
     if !NNlib.flipkernel(cdims)
         set_mlir_data!(
-            dw, get_mlir_data(Reactant.Ops.reverse(dw; dimensions=output_spatial_dims))
+            dw, get_mlir_data(@opcall(reverse(dw; dimensions=output_spatial_dims)))
         )
     end
 
@@ -312,7 +312,7 @@ function overloaded_∇conv_data!(
     )
 
     if NNlib.flipkernel(cdims)
-        w = Reactant.Ops.reverse(w; dimensions=kernel_spatial_dims)
+        w = @opcall reverse(w; dimensions=kernel_spatial_dims)
     end
 
     conv = MLIR.Dialects.stablehlo.convolution(
@@ -411,7 +411,7 @@ function NNlib.batched_mul!(
         end
     end
 
-    tmp = Ops.dot_general(
+    tmp = @opcall dot_general(
         T1.(materialize_traced_array(x)),
         T1.(materialize_traced_array(y));
         contracting_dimensions=([3], [2]),
@@ -430,7 +430,7 @@ function NNlib.pad_constant(
     low = [i[1] for i in pad]
     high = [i[2] for i in pad]
     interior = [0 for i in pad]
-    return Ops.pad(materialize_traced_array(x), value; low, high, interior)
+    return @opcall pad(materialize_traced_array(x), value; low, high, interior)
 end
 
 # Gather
@@ -462,7 +462,7 @@ end
 function _nnlib_gather_impl(src::AnyTracedRArray, idxs::AbstractArray, n_dims::Int)
     idxs = TracedUtils.promote_to(TracedRArray{Int,ndims(idxs)}, idxs)
     n_idxs = size(idxs, 1)
-    return Ops.gather(
+    return @opcall gather(
         src,
         idxs;
         offset_dims=collect(Int64, 1:n_dims),
@@ -506,7 +506,7 @@ function NNlib.scatter(
         )
     end
     xinit = isnothing(init) ? NNlib.scatter_empty(op, T) : init
-    dst = Ops.fill(xinit, dstsz)
+    dst = @opcall fill(xinit, dstsz)
 
     NNlib.scatter!(op, dst, src, idx)
     return dst
@@ -551,17 +551,21 @@ function _nnlib_scatter_impl(
 ) where {OP,T}
     scatter_indices = TracedUtils.promote_to(TracedRArray{Int,ndims(idx)}, idx)
     n_idxs = size(scatter_indices, 1)
-    return Ops.scatter(
-        op,
-        [dst],
-        scatter_indices,
-        [src];
-        update_window_dims=collect(Int64, 1:n_dims),
-        inserted_window_dims=collect(Int64, (n_dims + 1):ndims(dst)),
-        input_batching_dims=Int64[],
-        scatter_indices_batching_dims=Int64[],
-        scatter_dims_to_operand_dims=collect(Int64, (ndims(dst) - n_idxs + 1):ndims(dst)),
-        index_vector_dim=Int64(1),
+    return @opcall(
+        scatter(
+            op,
+            [dst],
+            scatter_indices,
+            [src];
+            update_window_dims=collect(Int64, 1:n_dims),
+            inserted_window_dims=collect(Int64, (n_dims + 1):ndims(dst)),
+            input_batching_dims=Int64[],
+            scatter_indices_batching_dims=Int64[],
+            scatter_dims_to_operand_dims=collect(
+                Int64, (ndims(dst) - n_idxs + 1):ndims(dst)
+            ),
+            index_vector_dim=Int64(1),
+        )
     )[1]
 end
 

ext/ReactantNNlibExt/Ops.jl

@@ -12,17 +12,19 @@ function reduce_window(
 
     padding = collect(Int64, reshape([padding..., 0, 0, 0, 0], (2, N))')
 
-    return Ops.reduce_window(
-        f,
-        [materialize_traced_array(x)],
-        [Ops.constant(T(init))];
-        window_dimensions=[kernel_size..., 1, 1],
-        window_strides=[stride..., 1, 1],
-        window_dilations=[dilation..., 1, 1],
-        padding_low=padding[:, 1],
-        padding_high=padding[:, 2],
-        output_shape=Int[output_spatial_shapes..., size(x, N - 1), size(x, N)],
-        base_dilations=ones(Int, N),
+    return @opcall(
+        reduce_window(
+            f,
+            [materialize_traced_array(x)],
+            [@opcall(constant(T(init)))];
+            window_dimensions=[kernel_size..., 1, 1],
+            window_strides=[stride..., 1, 1],
+            window_dilations=[dilation..., 1, 1],
+            padding_low=padding[:, 1],
+            padding_high=padding[:, 2],
+            output_shape=Int[output_spatial_shapes..., size(x, N - 1), size(x, N)],
+            base_dilations=ones(Int, N),
+        )
     )[1]
 end
 
@@ -31,7 +33,7 @@ function upsample_linear(
 ) where {T}
     W, _, _ = size(x)
 
-    out_idxs = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
+    out_idxs = @opcall iota(Int32, [out_size[1]]; iota_dimension=1)
     iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_idxs, align_corners, W)
 
     x0 = x[iw0, :, :]
@@ -45,8 +47,8 @@ function upsample_linear(
 ) where {T}
     W, H, _, _ = size(x)
 
-    out_width = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
-    out_height = Ops.iota(Int32, [out_size[2]]; iota_dimension=1)
+    out_width = @opcall iota(Int32, [out_size[1]]; iota_dimension=1)
+    out_height = @opcall iota(Int32, [out_size[2]]; iota_dimension=1)
 
     iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_width, align_corners, W)
     ih0, ih1, h0_λ, h1_λ = source_idx_and_λ(rheight, out_height, align_corners, H)
@@ -72,9 +74,9 @@ function upsample_linear(
 ) where {T}
     W, H, D, _, _ = size(x)
 
-    out_width = Ops.iota(Int32, [out_size[1]]; iota_dimension=1)
-    out_height = Ops.iota(Int32, [out_size[2]]; iota_dimension=1)
-    out_depth = Ops.iota(Int32, [out_size[3]]; iota_dimension=1)
+    out_width = @opcall iota(Int32, [out_size[1]]; iota_dimension=1)
+    out_height = @opcall iota(Int32, [out_size[2]]; iota_dimension=1)
+    out_depth = @opcall iota(Int32, [out_size[3]]; iota_dimension=1)
 
     iw0, iw1, w0_λ, w1_λ = source_idx_and_λ(rwidth, out_width, align_corners, W)
     ih0, ih1, h0_λ, h1_λ = source_idx_and_λ(rheight, out_height, align_corners, H)

ext/ReactantNNlibExt/ReactantNNlibExt.jl

@@ -7,6 +7,7 @@ using Reactant:
 
 using Reactant.TracedUtils:
     TracedUtils, materialize_traced_array, get_mlir_data, set_mlir_data!
+using Reactant.Ops: @opcall
 
 using ReactantCore: @trace
 using LinearAlgebra: LinearAlgebra, triu

ext/ReactantOneHotArraysExt.jl

@@ -3,6 +3,7 @@ module ReactantOneHotArraysExt
 using OneHotArrays
 using Reactant
 using Reactant: TracedRArray, TracedRNumber, TracedUtils, Ops
+using Reactant.Ops: @opcall
 
 function Reactant.traced_type_inner(
     @nospecialize(_::Type{OneHotArrays.OneHotArray{T,N,Np1,I}}),
@@ -35,9 +36,9 @@ function TracedUtils.materialize_traced_array(r::OneHotArrays.OneHotArray)
 
     linear_indices =
         TracedUtils.promote_to(TracedRArray{Int64,ndims(r.indices)}, indices) .+
-        Ops.iota(Int64, [B]; iota_dimension=1) .* N
+        @opcall(iota(Int64, [B]; iota_dimension=1)) .* N
 
-    z = Ops.fill(false, (N, B))
+    z = @opcall(fill(false, (N, B)))
     z[linear_indices] = fill(true, length(linear_indices))
     return reshape(z, size(r))
 end

ext/ReactantPythonCallExt/ReactantPythonCallExt.jl

@@ -2,6 +2,7 @@ module ReactantPythonCallExt
 
 using PythonCall
 using Reactant: Reactant, TracedRArray
+using Reactant.Ops: @opcall
 
 const jaxptr = Ref{Py}()
 const jnpptr = Ref{Py}()

ext/ReactantPythonCallExt/pycall.jl

@@ -12,7 +12,7 @@ function PythonCall.pycall(f::Py, arg0::TracedRArray, argNs::TracedRArray...; kw
     end
 
     lowered = jax.jit(f).lower(inputs...)
-    res = Reactant.Ops.hlo_call(pyconvert(String, lowered.as_text()), arg0, argNs...)
+    res = @opcall hlo_call(pyconvert(String, lowered.as_text()), arg0, argNs...)
 
     return length(res) == 0 ? nothing : res[1]
 end

ext/ReactantSpecialFunctionsExt.jl

@@ -2,16 +2,17 @@ module ReactantSpecialFunctionsExt
 using SpecialFunctions
 using Reactant: Ops, Reactant, TracedRNumber, ReactantFloat, ReactantInt, ReactantFloatInt
 using Reactant.TracedRNumberOverrides: float
+using Reactant.Ops: @opcall
 
 for fn in [:digamma, :erf, :erfc, (:loggamma, :lgamma)]
     (fns, fno) = fn isa Tuple ? fn : (fn, fn)
     @eval(function SpecialFunctions.$fns(x::TracedRNumber{<:ReactantFloatInt})
-        return Ops.$fno(float(x))
+        return @opcall $fno(float(x))
     end)
 end
 
 function SpecialFunctions.gamma(x::TracedRNumber{<:ReactantFloat})
-    return exp(Ops.lgamma(float(x)))
+    return exp(@opcall(lgamma(float(x))))
 end
 
 function SpecialFunctions.gamma(n::TracedRNumber{<:ReactantInt})
@@ -29,13 +30,14 @@ end
 # SpecialFunctions.invdigamma
 
 function SpecialFunctions.trigamma(x::TracedRNumber{<:ReactantFloatInt})
-    return Ops.polygamma(Ops.constant(Float64(1)), float(x))#TODO: change Ops definition
+    #TODO: change Ops definition
+    return @opcall(polygamma(@opcall(constant(Float64(1))), float(x)))
 end
 
 function SpecialFunctions.polygamma(
     n::TracedRNumber{<:ReactantFloatInt}, x::TracedRNumber{<:ReactantFloatInt}
 )
-    return Ops.polygamma(float(n), float(x))
+    return @opcall polygamma(float(n), float(x))
 end
 
 # SpecialFunctions.gamma_inc
@@ -112,7 +114,7 @@ end
 function SpecialFunctions.zeta(
     z::TracedRNumber{T}, s::TracedRNumber{T}
 ) where {T<:ReactantFloatInt}
-    return Ops.zeta(z, s)
+    return @opcall zeta(z, s)
 end
 
 end # module ReactantSpecialFunctionsExt

src/ConcreteRArray.jl

@@ -750,7 +750,7 @@ function Base.fill(
 )
     output_shardings = Sharding.is_sharded(sharding) ? Dict(1 => sharding) : nothing
     fn = Reactant.compile((); output_shardings) do
-        return Ops.fill(val, collect(Int64, dims))
+        return @opcall fill(val, collect(Int64, dims))
     end
     return fn()
 end