add support for complex tensor types in getConstantAttr

[mofeing]

diff rules that use HLOConstantFP in Enzyme-JAX are broken for complex numbers, due to mlir::enzyme::getConstantAttr not being able to handle them.

for example, the diff rule of stablehlo.rsqrt is defined as the following:

def : HLODerivative<"RsqrtOp", (Op $x), [
    (CheckedDiv (DiffeRet), (Mul (HLOConstantFP<"-2"> $x), (Mul $x, (Sqrt $x))))
]>;

which works for real numbers. but for complex numbers, like the following example,

func.func @rsqrt_complex(%x : tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>> {
  %y = stablehlo.rsqrt %x : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
  func.return %y : tensor<2xcomplex<f32>>
}

it errors:

mofeing@hydra:~/Enzyme-JAX$ ./bazel-bin/enzymexlamlir-opt --enzyme-wrap="infn=rsqrt_complex outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse test/lit_tests/diffrules/stablehlo/rsqrt.mlir 
 unsupported eltype: <<NULL TYPE>> of type tensor<2xcomplex<f32>>

we haven't checked this because we don't test complex numbers on diff rules that do not have a specific behaviour for them (i.e. SelectIfComplex).

this pr adds support for complex numbers in getConstantAttr. the current semantics is to create a ComplexAttr or a DenseElementsAttr with ComplexType eltype with the given value as the real part. not sure if we want to write a way to set the imaginary part.

[wsmoses]

can you add a relevant test?

[wsmoses]

format fails, otherwise all is good

[wsmoses]

though also what happened to the test?

[mofeing]

I removed it because complex::ConstantOp doesn't accept a TypedAttr and we can't simply cast to ArrayAttr, which is what it accepts.

I'm gonna submit now a test for math::RsqrtOp which should work with complex tensors, and with that, we should be able to merge when the test passes.

[mofeing]

okay, just found out that the diff rule test of arith.select should be enough.

funny how arith.constant can be used to instantiate a complex tensor value, but not a complex value 🫠

[wsmoses]

[0/2] Re-checking globbed directories...
[0/1] Running MLIR regression tests
-- Testing: 129 tests, 32 workers --
Testing:  0.. 10.. 20.. 30.. 40.. 50..
FAIL: Enzyme :: MLIR/ReverseMode/arith.mlir (78 of 129)
******************** TEST 'Enzyme :: MLIR/ReverseMode/arith.mlir' FAILED ********************
Exit Code: 1

Command Output (stderr):
--
/__w/Enzyme/Enzyme/Enzyme/enzyme-build/Enzyme/MLIR/enzymemlir-opt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir | FileCheck /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir # RUN: at line 1
+ /__w/Enzyme/Enzyme/Enzyme/enzyme-build/Enzyme/MLIR/enzymemlir-opt --split-input-file --enzyme --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir
+ FileCheck /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir
/__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir:27:10: error: 'complex.add' op operand #0 must be complex type with floating-point elements, but got 'tensor<complex<f64>>'
  %res = arith.select %c, %a, %b : tensor<complex<f64>>
         ^
/__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir:27:10: note: see current operation: %9 = "complex.add"(%8, %arg3) <{fastmath = #arith.fastmath<none>}> : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
/__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir:40:11: error: CHECK: expected string not found in input
// CHECK: func.func private @diffeselect_tensor_complex(%[[c:.+]]: i1, %[[a:.+]]: tensor<complex<f64>>, %[[b:.+]]: tensor<complex<f64>>, %[[dr:.+]]: tensor<complex<f64>>) -> (tensor<complex<f64>>, tensor<complex<f64>>) {
          ^
<stdin>:15:3: note: scanning from here
 }
  ^

Input file: <stdin>
Check file: /__w/Enzyme/Enzyme/Enzyme/enzyme/test/MLIR/ReverseMode/arith.mlir

-dump-input=help explains the following input dump.

Input was:
<<<<<<
          .
          .
          .
         10:  func.func private @diffeselect(%arg0: i1, %arg1: f64, %arg2: f64, %arg3: f64) -> (f64, f64) { 
         11:  %cst = arith.constant 0.000000e+00 : f64 
         12:  %0 = arith.select %arg0, %arg3, %cst : f64 
         13:  %1 = arith.select %arg0, %cst, %arg3 : f64 
         14:  return %0, %1 : f64, f64 
         15:  } 
check:40       X error: no match found
         16: } 
check:40     ~~
         17:  
check:40     ~
         18: // ----- 
check:40     ~~~~~~~~~
>>>>>>

--

********************
Testing:  0.. 10.. 20.. 30.. 40.. 50.. 60.. 70.. 80.. 90.. 
********************
Failed Tests (1):
  Enzyme :: MLIR/ReverseMode/arith.mlir


Testing Time: 0.41s

Total Discovered Tests: 129
  Passed           : 127 (98.45%)
  Expectedly Failed:   1 (0.78%)
  Failed           :   1 (0.78%)
FAILED: test/MLIR/CMakeFiles/check-enzymemlir /__w/Enzyme/Enzyme/Enzyme/enzyme-build/test/MLIR/CMakeFiles/check-enzymemlir 

[mofeing]

I am aware. The problem here is that complex ops do not act element-wise and do not accept tensor args.

I've tried with math.sqrt and with the forward diff of arith.divf too but math and arith ops args have a FloatLike type constraint that forbids passing complex and tensor complex values.

// -----
../test/MLIR/ReverseMode/math.mlir:27:10: error: 'math.sqrt' op operand #0 must be floating-point-like, but got 'tensor<complex<f64>>'
  %res = math.sqrt %x : tensor<complex<f64>>
         ^
../test/MLIR/ReverseMode/math.mlir:27:10: note: see current operation: %0 = "math.sqrt"(%arg0) <{fastmath = #arith.fastmath<none>}> : (tensor<complex<f64>>) -> tensor<complex<f64>>
// -----

Not even with --verify-each=0 I manage to remove that error.

It's weird that you can instantiate a complex tensor like this

%0 = arith.constant dense<(0.0,0.0)> : tensor<complex<f64>>

but you cannot use any arith nor math op on it.

I think it's untestable right now: it requires a discussion and fix upstream. But at least it seems to fix the problems in stablehlo as tests in EnzymeAD/Enzyme-JAX#2426 pass now.