Julia coding patterns that maps to XLA/scan resp. StableHLO/while?

[oschulz]

I'd like to accelerate custom IIR-DSP filters (biquad, but also more application-specific, in any case inherently serial along signal-samples but parallel across signals) via Reactant. From what I understand, the corrects resulting StableHLO-pattern should be a while loop along time over operations broadcasted across samples (e.g. like in RNNs are handled)?

As a test, I tried implementing a cumulative sum manually, but the way I wrote it, Reactant unrolls the loop:

julia> using Reactant; Reactant.set_default_backend("cpu")

julia> A = rand(Float32, 7, 5); rA = Reactant.to_rarray(A);

julia> function foo!(A)
           xs = eachrow(A)
           for i in eachindex(xs)[begin+1:end]
               xs[i] += xs[i-1]
           end
           return A
       end
foo! (generic function with 1 method)

julia> foo!(A)
7×5 Matrix{Float32}:
 0.385297  0.487759  0.478373  0.591901  0.940462
 1.33295   0.890929  0.481439  1.03631   0.970324
 1.96804   1.06544   1.42751   1.55177   1.84208
 2.96404   1.2932    2.05575   1.83932   2.3411
 3.18363   1.68394   2.55661   1.8492    3.05771
 4.15348   1.74053   3.53721   1.88457   3.54661
 4.59739   1.75737   3.95119   1.94483   4.04031

julia> @code_hlo foo!(rA)
module @"reactant_foo!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<5x7xf32> {tf.aliasing_output = 0 : i32}) -> tensor<5x7xf32> {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<5x7xf32>) -> tensor<7x5xf32>
    %1 = stablehlo.slice %0 [1:2, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %2 = stablehlo.slice %0 [0:1, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %3 = stablehlo.slice %0 [2:3, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %4 = stablehlo.slice %0 [3:4, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %5 = stablehlo.slice %0 [4:5, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %6 = stablehlo.slice %0 [5:6, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %7 = stablehlo.slice %0 [6:7, 0:5] : (tensor<7x5xf32>) -> tensor<1x5xf32>
    %8 = stablehlo.reshape %2 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %9 = stablehlo.reshape %1 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %10 = stablehlo.add %9, %8 : tensor<5x1xf32>
    %11 = stablehlo.reshape %3 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %12 = stablehlo.add %11, %10 : tensor<5x1xf32>
    %13 = stablehlo.reshape %4 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %14 = stablehlo.add %13, %12 : tensor<5x1xf32>
    %15 = stablehlo.reshape %5 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %16 = stablehlo.add %15, %14 : tensor<5x1xf32>
    %17 = stablehlo.reshape %6 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %18 = stablehlo.add %17, %16 : tensor<5x1xf32>
    %19 = stablehlo.reshape %7 : (tensor<1x5xf32>) -> tensor<5x1xf32>
    %20 = stablehlo.add %19, %18 : tensor<5x1xf32>
    %21 = stablehlo.concatenate %8, %10, %12, %14, %16, %18, %20, dim = 1 : (tensor<5x1xf32>, tensor<5x1xf32>, tensor<5x1xf32>, tensor<5x1xf32>, tensor<5x1xf32>, tensor<5x1xf32>, tensor<5x1xf32>) -> tensor<5x7xf32>
    return %21 : tensor<5x7xf32>
  }
}

Obviously this won't scale to larger arrays (more signal time steps). Should Reactant do this automatically, or do I need to give it some hints somehow?

[wsmoses]

You should write @trace for and that should do it.

There is some ongoing work to make it automatically do this upgrade but for now definitely add the trace yourself

[oschulz]

Thanks! With @trace we run into this error, in this case:

julia> A = rand(Float32, 7, 5); rA = Reactant.to_rarray(A);

julia> using Reactant; Reactant.set_default_backend("cpu")

julia> A = rand(Float32, 7, 5); rA = Reactant.to_rarray(A);

julia> function foo!(A)
           xs = eachrow(A)
           @trace for i in eachindex(xs)[begin+1:end]
               xs[i] += xs[i-1]
           end
           return A
       end
foo! (generic function with 1 method)

julia> foo!(A)
7×5 Matrix{Float32}: [...]

julia> @code_hlo foo!(rA)
┌ Warning: `Adapt.parent_type` is not implemented for Base.OneTo{Int64}. Assuming Base.OneTo{Int64} isn't a wrapped array.
└ @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Reactant.jl:50
ERROR: TypeError: in OneTo, in T, expected T<:Integer, got Type{Reactant.TracedRNumber{Int64}}

Do I need to do something with the iteration range as well?

[wsmoses]

cc @Pangoraw

[oschulz]

Ah, so this works:

julia> function foo!(A)
           @trace for i in axes(A,1)[2:end]
               A[i,:] += A[i-1,:]
           end
           return A
       end
foo! (generic function with 2 methods)

julia> @code_hlo foo!(rA)
module @"reactant_foo!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<5x7xf32> {tf.aliasing_output = 0 : i32}) -> tensor<5x7xf32> {
    %c = stablehlo.constant dense<0> : tensor<i32>
    %c_0 = stablehlo.constant dense<0> : tensor<i64>
    %c_1 = stablehlo.constant dense<1> : tensor<i64>
    %c_2 = stablehlo.constant dense<2> : tensor<i64>
    %c_3 = stablehlo.constant dense<6> : tensor<i64>
    %0:2 = stablehlo.while(%iterArg = %c_0, %iterArg_4 = %arg0) : tensor<i64>, tensor<5x7xf32> attributes {enzymexla.disable_min_cut}
    cond {
      %1 = stablehlo.compare  LT, %iterArg, %c_3 : (tensor<i64>, tensor<i64>) -> tensor<i1>
      stablehlo.return %1 : tensor<i1>
    } do {
      %1 = stablehlo.transpose %iterArg_4, dims = [1, 0] : (tensor<5x7xf32>) -> tensor<7x5xf32>
      %2 = stablehlo.add %c_2, %iterArg : tensor<i64>
      %3 = stablehlo.add %iterArg, %c_1 : tensor<i64>
      %4 = stablehlo.subtract %2, %c_1 : tensor<i64>
      %5 = stablehlo.convert %4 : (tensor<i64>) -> tensor<i32>
      %6 = stablehlo.dynamic_slice %1, %5, %c, sizes = [1, 5] : (tensor<7x5xf32>, tensor<i32>, tensor<i32>) -> tensor<1x5xf32>
      %7 = stablehlo.convert %iterArg : (tensor<i64>) -> tensor<i32>
      %8 = stablehlo.dynamic_slice %1, %7, %c, sizes = [1, 5] : (tensor<7x5xf32>, tensor<i32>, tensor<i32>) -> tensor<1x5xf32>
      %9 = stablehlo.add %6, %8 : tensor<1x5xf32>
      %10 = stablehlo.reshape %9 : (tensor<1x5xf32>) -> tensor<5x1xf32>
      %11 = stablehlo.dynamic_update_slice %iterArg_4, %10, %c, %5 : (tensor<5x7xf32>, tensor<5x1xf32>, tensor<i32>, tensor<i32>) -> tensor<5x7xf32>
      stablehlo.return %3, %11 : tensor<i64>, tensor<5x7xf32>
    }
    return %0#1 : tensor<5x7xf32>
  }
}

But this doesn't (would be necessary for actual digital filters, I think):

julia> function foo!(A)
           @trace for i in axes(A,1)[2:end]
               A[i,:] .+= A[i-1,:]
           end
           return A
       end
foo! (generic function with 2 methods)

julia> @code_hlo foo!(rA)
ERROR: ArgumentError: unable to check bounds for indices of type Reactant.TracedRNumber{Int64}
Stacktrace:
  [1] checkindex(::Type{Bool}, inds::Base.OneTo{Int64}, i::Reactant.TracedRNumber{Int64})
    @ Base ./abstractarray.jl:751
  [2] checkbounds_indices
    @ ./abstractarray.jl:725 [inlined]

Using eachrow doesn't work due to the Base.OneTo problem, but it seems related to the Slices object, not the iteration range:

ulia> function foo!(A)
           xs = eachrow(A)
           @trace for i in axes(A,1)[2:end]
               xs[i] += xs[i-1]
           end
           return A
       end
foo! (generic function with 2 methods)

julia> @code_hlo foo!(rA)
ERROR: TypeError: in OneTo, in T, expected T<:Integer, got Type{Reactant.TracedRNumber{Int64}}
Stacktrace:
  [1] traced_type_inner(T::Type, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:734
  [2] traced_tuple_type_inner(T::Type{…}, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:146

[wsmoses]

glad you figured a workaround, we still should fix the issues referred to above, so reopening

[wsmoses]

also can you include the full backtrace for all of the above

[oschulz]

Thanks, sure (I'll do a separate post for each). Sorry, I closed the issue by mistake. Here with broadcasted +=:

julia> function foo!(A)
           @trace for i in axes(A,1)[2:end]
               A[i,:] .+= A[i-1,:]
           end
           return A
       end
foo! (generic function with 2 methods)

julia> @code_hlo foo!(rA)
ERROR: ArgumentError: unable to check bounds for indices of type Reactant.TracedRNumber{Int64}
Stacktrace:
  [1] checkindex(::Type{Bool}, inds::Base.OneTo{Int64}, i::Reactant.TracedRNumber{Int64})
    @ Base ./abstractarray.jl:751
  [2] checkbounds_indices
    @ ./abstractarray.jl:725 [inlined]
  [3] checkbounds
    @ ./abstractarray.jl:681 [inlined]
  [4] checkbounds
    @ ./abstractarray.jl:699 [inlined]
  [5] view
    @ ./subarray.jl:214 [inlined]
  [6] maybeview
    @ ./views.jl:149 [inlined]
  [7] dotview
    @ ./broadcast.jl:1228 [inlined]
  [8] macro expansion
    @ ./REPL[53]:3 [inlined]
  [9] #76
    @ /user/.julia/packages/ReactantCore/9hY4Z/src/ReactantCore.jl:352 [inlined]
 [10] (::Nothing)(none::var"#76#78", none::Base.RefValue{…}, none::Base.RefValue{…}, none::Base.RefValue{…}, none::Base.RefValue{…}, none::Base.RefValue{…})
    @ Reactant ./<missing>:0
 [11] getproperty
    @ ./Base.jl:49 [inlined]
 [12] getindex
    @ ./refvalue.jl:59 [inlined]
 [13] #76
    @ /user/.julia/packages/ReactantCore/9hY4Z/src/ReactantCore.jl:248 [inlined]
 [14] call_with_reactant(::var"#76#78", ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…})
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/utils.jl:0
 [15] make_mlir_fn(f::var"#76#78", args::Vector{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Nothing, verify_arg_names::NTuple{…}, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils /user/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:332
 [16] make_mlir_fn
    @ /user/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:262 [inlined]
 [17] while_loop(::var"#75#77", ::var"#76#78", ::Base.RefValue{…}, ::Vararg{…}; track_numbers::Type, verify_arg_names::NTuple{…}, checkpointing::Bool, mincut::Bool, location::Reactant.MLIR.IR.Location)
    @ Reactant.Ops /user/.julia/packages/Reactant/gBXlB/src/Ops.jl:2067
 [18] #traced_while#127
    @ /user/.julia/packages/Reactant/gBXlB/src/ControlFlow.jl:20 [inlined]
 [19] traced_while
    @ /user/.julia/packages/Reactant/gBXlB/src/ControlFlow.jl:11 [inlined]
 [20] macro expansion
    @ /user/.julia/packages/ReactantCore/9hY4Z/src/ReactantCore.jl:260 [inlined]
 [21] foo!
    @ ./REPL[53]:2 [inlined]
 [22] (::Nothing)(none::typeof(foo!), none::Reactant.TracedRArray{Float32, 2})
    @ Reactant ./<missing>:0
 [23] getproperty
    @ ./Base.jl:49 [inlined]
 [24] size
    @ /user/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:527 [inlined]
 [25] axes
    @ ./abstractarray.jl:98 [inlined]
 [26] axes
    @ ./abstractarray.jl:77 [inlined]
 [27] macro expansion
    @ /user/.julia/packages/ReactantCore/9hY4Z/src/ReactantCore.jl:316 [inlined]
 [28] foo!
    @ ./REPL[53]:2 [inlined]
 [29] call_with_reactant(::typeof(foo!), ::Reactant.TracedRArray{Float32, 2})
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/utils.jl:0
 [30] make_mlir_fn(f::typeof(foo!), args::Tuple{…}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils /user/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:332
 [31] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1555
 [32] compile_mlir!
    @ /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1522 [inlined]
 [33] (::Reactant.Compiler.var"#9#10"{Nothing, typeof(foo!), Tuple{ConcretePJRTArray{…}}})(ctx::Reactant.MLIR.IR.Context)
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1362
 [34] with_context(f::Reactant.Compiler.var"#9#10"{Nothing, typeof(foo!), Tuple{…}}; allow_use_existing::Bool)
    @ Reactant.MLIR.IR /user/.julia/packages/Reactant/gBXlB/src/mlir/IR/Context.jl:99
 [35] with_context(f::Function)
    @ Reactant.MLIR.IR /user/.julia/packages/Reactant/gBXlB/src/mlir/IR/Context.jl:82
 [36] compile_mlir(f::Function, args::Tuple{…}; client::Nothing, kwargs::@Kwargs{…})
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1358
 [37] top-level scope
    @ /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:2586

I tried using @inbounds, but that didn't seem to help.

[oschulz]

Here with eachrow (still using for i in axes(A,1)[2:end] for minimal changes):

julia> function foo!(A)
           xs = eachrow(A)
           @trace for i in axes(A,1)[2:end]
               xs[i] += xs[i-1]
           end
           return A
       end
foo! (generic function with 2 methods)

julia> @code_hlo foo!(rA)
ERROR: TypeError: in OneTo, in T, expected T<:Integer, got Type{Reactant.TracedRNumber{Int64}}
Stacktrace:
  [1] traced_type_inner(T::Type, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:734
  [2] traced_tuple_type_inner(T::Type{…}, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:146
  [3] traced_type_inner(T::Type, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:625
  [4] traced_type_inner(T::Type, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:676
  [5] traced_type_inner(PT::Type{…}, seen::Dict{…}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:565
  [6] traced_type(T::Type, ::Val{…}, track_numbers::Type, sharding::Reactant.Sharding.NoSharding, runtime::Nothing)
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:870
  [7] make_tracer_unknown(seen::Reactant.OrderedIdDict{…}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{})
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:1032
  [8] make_tracer_unknown
    @ /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:1009 [inlined]
  [9] #make_tracer#131
    @ /user/.julia/packages/Reactant/gBXlB/src/Tracing.jl:1146 [inlined]
 [10] while_loop(::var"#79#81", ::var"#80#82", ::Base.RefValue{…}, ::Vararg{…}; track_numbers::Type, verify_arg_names::NTuple{…}, checkpointing::Bool, mincut::Bool, location::Reactant.MLIR.IR.Location)
    @ Reactant.Ops /user/.julia/packages/Reactant/gBXlB/src/Ops.jl:2039
 [11] #traced_while#127
    @ /user/.julia/packages/Reactant/gBXlB/src/ControlFlow.jl:20 [inlined]
 [12] traced_while
    @ /user/.julia/packages/Reactant/gBXlB/src/ControlFlow.jl:11 [inlined]
 [13] macro expansion
    @ /user/.julia/packages/ReactantCore/9hY4Z/src/ReactantCore.jl:260 [inlined]
 [14] foo!
    @ ./REPL[55]:3 [inlined]
 [15] (::Nothing)(none::typeof(foo!), none::Reactant.TracedRArray{Float32, 2})
    @ Reactant ./<missing>:0
 [16] getproperty
    @ ./Base.jl:49 [inlined]
 [17] size
    @ /user/.julia/packages/Reactant/gBXlB/src/TracedRArray.jl:527 [inlined]
 [18] axes
    @ ./abstractarray.jl:98 [inlined]
 [19] axes
    @ ./abstractarray.jl:77 [inlined]
 [20] #326
    @ ./slicearray.jl:61 [inlined]
 [21] map
    @ ./tuple.jl:355 [inlined]
 [22] _eachslice
    @ ./slicearray.jl:61 [inlined]
 [23] eachrow
    @ ./slicearray.jl:166 [inlined]
 [24] foo!
    @ ./REPL[55]:2 [inlined]
 [25] call_with_reactant(::typeof(foo!), ::Reactant.TracedRArray{Float32, 2})
    @ Reactant /user/.julia/packages/Reactant/gBXlB/src/utils.jl:0
 [26] make_mlir_fn(f::typeof(foo!), args::Tuple{…}, kwargs::@NamedTuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils /user/.julia/packages/Reactant/gBXlB/src/TracedUtils.jl:332
 [27] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1555
 [28] compile_mlir!
    @ /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1522 [inlined]
 [29] (::Reactant.Compiler.var"#9#10"{Nothing, typeof(foo!), Tuple{ConcretePJRTArray{…}}})(ctx::Reactant.MLIR.IR.Context)
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1362
 [30] with_context(f::Reactant.Compiler.var"#9#10"{Nothing, typeof(foo!), Tuple{…}}; allow_use_existing::Bool)
    @ Reactant.MLIR.IR /user/.julia/packages/Reactant/gBXlB/src/mlir/IR/Context.jl:99
 [31] with_context(f::Function)
    @ Reactant.MLIR.IR /user/.julia/packages/Reactant/gBXlB/src/mlir/IR/Context.jl:82
 [32] compile_mlir(f::Function, args::Tuple{…}; client::Nothing, kwargs::@Kwargs{…})
    @ Reactant.Compiler /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:1358
 [33] top-level scope
    @ /user/.julia/packages/Reactant/gBXlB/src/Compiler.jl:2586