sqrt(::Complex) unsupported due to conversion exceptions

[kmp5VT]

Hello,

I am trying to compute the square root of a matrix with values of ComplexF32 and I run into a internal Metal.jl error. Here is a small example

julia> using Metal, Test
julia> d = randn(ComplexF32, (2,2))
julia> md = mtl(d)
julia> dd = sqrt.(d)
julia> mdd = sqrt.(md)

This is the error I am seeing

ERROR: InvalidIRError: compiling MethodInstance for (::Metal.var"#broadcast_2d#203")(::MtlDeviceMatrix{…}, ::Base.Broadcast.Broadcasted{…}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to gpu_malloc)
Stacktrace:
 [1] malloc
   @ ~/.julia/packages/GPUCompiler/kqxyC/src/runtime.jl:88
 [2] macro expansion
   @ ~/.julia/packages/GPUCompiler/kqxyC/src/runtime.jl:183
 [3] macro expansion
   @ ./none:0
 [4] box
   @ ./none:0
 [5] box_float32
   @ ~/.julia/packages/GPUCompiler/kqxyC/src/runtime.jl:212
 [6] multiple call sites
   @ unknown:0

And this is my version info

pkg> st
  [dde4c033] Metal v1.1.0

julia> versioninfo()
Julia Version 1.10.4
Commit 48d4fd48430 (2024-06-04 10:41 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 10 × Apple M1 Max
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)
Environment:
  LD_LIBRARY_PATH = /Users/kpierce/Software/triqs/triqs_install/lib::/opt/intel/oneapi/mkl/latest/lib
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

Thanks!

[maleadt]

Dup of #69. You need to avoid exceptions in your code.

[mtfishman]

@maleadt to clarify a bit, is the issue that sqrt(::Complex) can potentially throw an exception, and that isn't allowed in broadcasting in Metal.jl right now? I don't see any exception throwing in sqrt(::Complex): https://github.com/JuliaLang/julia/blob/6f39acb9fb2b34b43d6b455d196b40f1759c9031/base/complex.jl#L530. I can't even guess from a mathematical perspective why it might throw an error, it seems like square root would be defined for any complex number but I guess there is exception throwing farther down the call stack in one of the math functions being called inside of sqrt(::Complex).

Would something actionable here be to add a @device_override definition for sqrt(::Complex), as a stopgap until exceptions are supported in Metal.jl? If so I would request to keep this open so we can track it and maybe try to fix this case, since it would be quite useful for us.

[maleadt]

is the issue that sqrt(::Complex) can potentially throw an exception, and that isn't allowed in broadcasting in Metal.jl right now?

Yeah, in fact, exceptions aren't allowed at all, not only in broadcasting code.

julia> function f(x)
       @inbounds x[1] = sqrt(x[1])
       return
       end
f (generic function with 1 method)

julia> @metal f(mtl([1f0+1f0im]))
ERROR: InvalidIRError: compiling MethodInstance for f(::MtlDeviceVector{ComplexF32, 1}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to gpu_malloc)

I don't see any exception throwing in sqrt(::Complex)

If you inspect the generated LLVM IR you can spot the calls to ijl_box_float32 as mentioned above. sqrt(::Complex) calls ssqs which contains a check conversion to Int.

Would something actionable here be to add a @device_override definition for sqrt(::Complex), as a stopgap until exceptions are supported in Metal.jl?

Sure, that, or a quirk for ssqs, throwing a GPU-compatible exception instead.