Fail to raise error with trig operations on OffsetArrays

[jlk9]

Running this script produces this error:

> grid_type=gaussian_islands XLA_FLAGS="--xla_dump_to=xla_dump" JULIA_NUM_THREADS=4 julia --color=yes --project -O0 simulations/ocean_climate_simulation_run.jl
[ Info: Oceananigans will use 4 threads
[ Info: Generating model...
[ Info: Compiling...
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1770149301.257092 3391981 pjrt_client.cc:551] PjRt-IFRT device count: total=1, addressable=1
I0000 00:00:1770149301.257121 3391981 pjrt_client.cc:555] Addressable PjRt-IFRT device: CpuDevice(id=0)
failed to raise func: func.func private @"##call__Z46gpu__my_interpolate_primary_atmospheric_state_16CompilerMetadataI16OffsetStaticSizeI12_0_97__0_49_E12DynamicCheckvv7NDRangeILi2E10StaticSizeI6_7__4_ES4_I8_16__16_E5TupleI5Int64S8_ES0_I8__1___1_EEE11OffsetArrayI7Float64Li2E13CuTracedArrayISE_Li2ELi1E9_112__64_EE#239$par0"(%arg0: memref<64x112xf64, 1>) {
  %cst = arith.constant 0.017453292519943295 : f64
  %cst_0 = arith.constant 2.000000e+03 : f64
  %cst_1 = arith.constant 2.000000e+00 : f64
  %cst_2 = arith.constant 1.000000e+03 : f64
  %cst_3 = arith.constant 57.295779513082323 : f64
  affine.parallel (%arg1, %arg2) = (0, 0) to (50, 98) {
    %0 = affine.load %arg0[%arg1 + 7, %arg2 + 7] : memref<64x112xf64, 1>
    %1 = affine.load %arg0[%arg1 + 7, %arg2 + 8] : memref<64x112xf64, 1>
    %2 = affine.load %arg0[%arg1 + 8, %arg2 + 7] : memref<64x112xf64, 1>
    %3 = affine.load %arg0[%arg1 + 8, %arg2 + 8] : memref<64x112xf64, 1>
    %4 = arith.subf %3, %1 {fastmathFlags = #llvm.fastmath<none>} : f64
    %5 = arith.mulf %4, %cst {fastmathFlags = #llvm.fastmath<none>} : f64
    %6 = arith.divf %5, %cst_0 {fastmathFlags = #llvm.fastmath<none>} : f64
    %7 = arith.subf %2, %0 {fastmathFlags = #llvm.fastmath<none>} : f64
    %8 = arith.mulf %7, %cst {fastmathFlags = #llvm.fastmath<none>} : f64
    %9 = arith.divf %8, %cst_0 {fastmathFlags = #llvm.fastmath<none>} : f64
    %10 = arith.addf %6, %9 {fastmathFlags = #llvm.fastmath<none>} : f64
    %11 = arith.divf %10, %cst_1 {fastmathFlags = #llvm.fastmath<none>} : f64
    %12 = arith.subf %3, %2 {fastmathFlags = #llvm.fastmath<none>} : f64
    %13 = arith.mulf %12, %cst {fastmathFlags = #llvm.fastmath<none>} : f64
    %14 = arith.divf %13, %cst_2 {fastmathFlags = #llvm.fastmath<none>} : f64
    %15 = arith.subf %1, %0 {fastmathFlags = #llvm.fastmath<none>} : f64
    %16 = arith.mulf %15, %cst {fastmathFlags = #llvm.fastmath<none>} : f64
    %17 = arith.divf %16, %cst_2 {fastmathFlags = #llvm.fastmath<none>} : f64
    %18 = arith.addf %14, %17 {fastmathFlags = #llvm.fastmath<none>} : f64
    %19 = arith.negf %18 {fastmathFlags = #llvm.fastmath<none>} : f64
    %20 = arith.divf %19, %cst_1 {fastmathFlags = #llvm.fastmath<none>} : f64
    %21 = arith.mulf %11, %11 {fastmathFlags = #llvm.fastmath<none>} : f64
    %22 = arith.mulf %20, %20 {fastmathFlags = #llvm.fastmath<none>} : f64
    %23 = arith.addf %21, %22 {fastmathFlags = #llvm.fastmath<none>} : f64
    %24 = math.sqrt %23 : f64
    %25 = arith.divf %11, %24 {fastmathFlags = #llvm.fastmath<none>} : f64
    %26 = arith.divf %20, %24 {fastmathFlags = #llvm.fastmath<none>} : f64
    %27 = arith.divf %26, %25 {fastmathFlags = #llvm.fastmath<none>} : f64
    %28 = math.atan %27 : f64
    %29 = arith.mulf %28, %cst_3 {fastmathFlags = #llvm.fastmath<none>} : f64
    %30 = math.isfinite %29 : f64
    scf.if %30 {
      llvm.intr.trap
    }
  }
  return
}
┌ Error: Compilation failed, MLIR module written to /var/folders/k4/1ytpy6y14y7_djhgb34ct0n80000gq/T/reactant_WcqUwO/module_000_nxc9_post_all_pm.mlir
└ @ Reactant.MLIR.IR ~/.julia/packages/Reactant/AyuUw/src/mlir/IR/Pass.jl:145
ERROR: LoadError: "failed to run pass manager on module"
Stacktrace:
 [1] run!(pm::Reactant.MLIR.IR.PassManager, operation::Reactant.MLIR.IR.Operation, key::String)
   @ Reactant.MLIR.IR ~/.julia/packages/Reactant/AyuUw/src/mlir/IR/Pass.jl:193
 [2] run_pass_pipeline!(mod::Reactant.MLIR.IR.Module, pass_pipeline::String, key::String; enable_verifier::Bool)
   @ Reactant.Compiler ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:1436
 [3] run_pass_pipeline!
   @ ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:1431 [inlined]
 [4] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(my_interpolate_state!), args::Tuple{ReactantState, OffsetMatrix{Float64, ConcreteIFRTArray{Float64, 2, Nothing}}}, compile_options::CompileOptions, callcache::Dict{Vector, @NamedTuple{f_name::String, mlir_result_types::Vector{Reactant.MLIR.IR.Type}, traced_result, mutated_args::Vector{Int64}, linear_results::Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, fnwrapped::Bool, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol}}, sdycache::Dict{Tuple{AbstractVector{Int64}, NTuple{var"#s1579", Symbol} where var"#s1579", NTuple{N, Int64} where N}, @NamedTuple{sym_name::Reactant.MLIR.IR.Attribute, mesh_attr::Reactant.MLIR.IR.Attribute, mesh_op::Reactant.MLIR.IR.Operation, mesh::Reactant.Sharding.Mesh}}, sdygroupidcache::Tuple{Reactant.Compiler.SdyGroupIDCounter{Int64}, IdDict{Union{Reactant.TracedRArray, Reactant.TracedRNumber}, Int64}}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{:IFRT}, legalize_stablehlo_to_mhlo::Bool, client::Reactant.XLA.IFRT.Client, kwargs::@Kwargs{})
   @ Reactant.Compiler ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:1919
 [5] compile_mlir!
   @ ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:1721 [inlined]
 [6] compile_xla(f::Function, args::Tuple{ReactantState, OffsetMatrix{Float64, ConcreteIFRTArray{Float64, 2, Nothing}}}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{compile_options::CompileOptions, fn_kwargs::@NamedTuple{}})
   @ Reactant.Compiler ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:3727
 [7] compile_xla
   @ ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:3700 [inlined]
 [8] compile(f::Function, args::Tuple{ReactantState, OffsetMatrix{Float64, ConcreteIFRTArray{Float64, 2, Nothing}}}; kwargs::@Kwargs{fn_kwargs::@NamedTuple{}, client::Nothing, reshape_propagate::Symbol, raise_first::Bool, assert_nonallocating::Bool, serializable::Bool, legalize_chlo_to_stablehlo::Bool, transpose_propagate::Symbol, donated_args::Symbol, optimize_then_pad::Bool, cudnn_hlo_optimize::Bool, compile_options::Missing, sync::Bool, no_nan::Bool, raise::Bool, shardy_passes::Symbol, optimize::Bool, optimize_communications::Bool})
   @ Reactant.Compiler ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:3815
 [9] top-level scope
   @ ~/.julia/packages/Reactant/AyuUw/src/Compiler.jl:2870
in expression starting at /Users/jkump/Desktop/NN_GB/GB-25/simulations/ocean_climate_simulation_run.jl:125

Will reduce further.

@wsmoses

[jlk9]

The above script has now eliminated Oceananigans as a dependency.

[wsmoses]

this......is weird

can you upload the pre_all mlir file?

[jlk9]

Here is the pre_all output:

mlir_output_error_0203.txt

[wsmoses]

okay the error here appears to come from the use of sind in a kernel, which is presently unsupported. it will be in time, but longer term [requires the custom overlay kernel MT @giordano I mentioned to you]

[giordano]

Do you mean #2131? I have no idea how to make that work.

[wsmoses]

yeah

[giordano]

That looks very hard to implement, because all the method table dispatch (including the one linked in #2131 (comment)) is based on the type of the configuration and of the job, but you want to purposefully reuse CUDA.jl's config and job, so short of somehow wrapping CUDA.jl jobs (which would be a terrible rabbit hole), I really don't see how to do that.