fft op crashes when being printing

[kuhar]

What happened?

The following snippet that uses the generic syntax crashes when printing:

  %0 = "stablehlo.fft"(%input) {
    fft_length = dense<32> : tensor<i64>, fft_type = #stablehlo<fft_type RFFT>
  } : (tensor<32xf32>) -> tensor<17xcomplex<f32>>

Error message:

$ tools/iree-compile --output-format=vm-bytecode --mlir-print-op-on-diagnostic=false --iree-hal-target-backends=vmvx --iree-input-type=stablehlo /usr/local/google/
home/kubak/iree/iree/tests/e2e/stablehlo_ops/fft.mlir -o check_vmvx_local-task_fft.mlir_module.vmfb --mlir-disable-threading --mlir-print-ir-after-all             
// -----// IR Dump After Canonicalizer (canonicalize) //----- //                                                                                                   
func.func @rfft_1d() {                                                                                                                                             
  %cst = arith.constant dense<[0.000000e+00, 0.000000e+00, -23.9563732, 0.000000e+00, -10.2543259, 0.000000e+00, -6.14436531, 0.000000e+00, -1.000000e+01, 0.000000
e+00, 3.86551499, 0.000000e+00, 0.637673854, 0.000000e+00, 52.4535065, 0.000000e+00, 0.000000e+00]> : tensor<17xf32>                                               
  %cst_0 = arith.constant dense<[6.668460e+02, 0.000000e+00, -590.16925, 0.000000e+00, 593.448486, 0.000000e+00, -579.528748, 0.000000e+00, 629.954041, 0.000000e+0
0, -567.11261, 0.000000e+00, 591.751465, 0.000000e+00, -583.189392, 0.000000e+00, 6.308460e+02]> : tensor<17xf32>                                                  
  %cst_1 = arith.constant dense<[9.000000e+00, 1.000000e+00, 4.500000e+00, -3.000000e-01, 1.000000e+01, -1.000000e+00, 5.500000e+00, 3.000000e-01, 2.990000e+02, 3.
500000e+00, -7.770000e-01, 2.000000e+00, 1.700000e+00, 3.500000e+00, -4.500000e+00, 0.000000e+00, 9.000000e+00, 1.000000e+00, 4.500000e+00, -3.000000e-01, 1.000000
e+01, -1.000000e+00, 5.500000e+00, 3.000000e-01, 2.990000e+02, 3.500000e+00, -7.770000e-01, 2.000000e+00, 1.700000e+00, 3.500000e+00, -4.500000e+00, 0.000000e+00]>
 : tensor<32xf32>                                                                                                                                                  
  %0 = util.optimization_barrier %cst_1 : tensor<32xf32>                                                                                                           
  %1 = stablehlo.fft %0, type =  RFFT, length = LLVM ERROR: printDenseI64Array only supports rank-1 arrays                                                         
Please report issues to https://github.com/openxla/iree/issues and include the crash backtrace.                                                                    
Stack dump:                                                                                                                                                        
0.      Program arguments: tools/iree-compile --output-format=vm-bytecode --mlir-print-op-on-diagnostic=false --iree-hal-target-backends=vmvx --iree-input-type=sta
blehlo /usr/local/google/home/kubak/iree/iree/tests/e2e/stablehlo_ops/fft.mlir -o check_vmvx_local-task_fft.mlir_module.vmfb --mlir-disable-threading --mlir-print-
ir-after-all                                                                                                                                                       
 #0 0x00007f83b2151f18 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/lib/Support/Unix/Si
gnals.inc:602:13                                                                                                                                                   
 #1 0x00007f83b2150180 llvm::sys::RunSignalHandlers() /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/lib/Support/Signals.cpp:105:18          
 #2 0x00007f83b2152578 SignalHandler(int) /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/lib/Support/Unix/Signals.inc:413:1                  
 #3 0x00007f83a865af90 (/lib/x86_64-linux-gnu/libc.so.6+0x3bf90)                                                                                                   
 #4 0x00007f83a86a9ccc __pthread_kill_implementation ./nptl/./nptl/pthread_kill.c:44:76                                                                            
 #5 0x00007f83a865aef2 raise ./signal/../sysdeps/posix/raise.c:27:6                                                                                                
 #6 0x00007f83a8645472 abort ./stdlib/./stdlib/abort.c:81:7                                                                                                        
 #7 0x00007f83b20c0365 llvm::report_fatal_error(llvm::Twine const&, bool) /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/lib/Support/ErrorHan
dling.cpp:125:5                                                                                                                                                    
 #8 0x00007f83b20c01a6 (/usr/local/google/home/kubak/iree/build-relass/lib/libIREECompiler.so.0+0x98c01a6)                                                         
 #9 0x00007f83acaeef9e auto mlir::DenseElementsAttr::getValues<long>() const /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/mlir/include/mlir/IR/B
uiltinAttributes.h:417:5                                                                                                                                           
#10 0x00007f83acaeef9e mlir::hlo::printDenseI64Array(mlir::OpAsmPrinter&, mlir::Operation*, mlir::DenseIntElementsAttr) /usr/local/google/home/kubak/iree/iree/thir
d_party/mlir-hlo/stablehlo/stablehlo/dialect/AssemblyFormat.cpp:260:38                                                                                             
#11 0x00007f83ac9e815a llvm::SmallVectorTemplateCommon<llvm::StringRef, void>::getFirstEl() const /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/l
lvm/include/llvm/ADT/SmallVector.h:133:46
#12 0x00007f83ac9e815a llvm::SmallVectorTemplateCommon<llvm::StringRef, void>::SmallVectorTemplateCommon(unsigned long) /usr/local/google/home/kubak/iree/iree/thir
d_party/llvm-project/llvm/include/llvm/ADT/SmallVector.h:138:49
#13 0x00007f83ac9e815a llvm::SmallVectorTemplateBase<llvm::StringRef, true>::SmallVectorTemplateBase(unsigned long) /usr/local/google/home/kubak/iree/iree/third_pa
rty/llvm-project/llvm/include/llvm/ADT/SmallVector.h:491:42
#14 0x00007f83ac9e815a llvm::SmallVectorImpl<llvm::StringRef>::SmallVectorImpl(unsigned int) /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/i
nclude/llvm/ADT/SmallVector.h:592:9
#15 0x00007f83ac9e815a llvm::SmallVector<llvm::StringRef, 2u>::SmallVector() /usr/local/google/home/kubak/iree/iree/third_party/llvm-project/llvm/include/llvm/ADT/
SmallVector.h:1202:19
#16 0x00007f83ac9e815a mlir::stablehlo::FftOp::print(mlir::OpAsmPrinter&) /usr/local/google/home/kubak/iree/build-relass/llvm-external-projects/mlir-hlo/stablehlo/
stablehlo/dialect/StablehloOps.cpp.inc:14490:45

This works after changing the length to tensor<1xi64>. I'd expect either a verifier error or this print work.

The equivalent code in the MHLO dialect prints fine: https://github.com/openxla/iree/blob/a8a70fb2dae68029fe7bdf159ed9353c7b531e1b/tests/e2e/xla_ops/fft.mlir

Steps to reproduce your issue

No response

Version information

mlir-hlo: git 65cc2df0392b3d0618fa70be8a10189760b7d30d, PiperOrigin-RevId: 534791162

[GleasonK]

Was this hand-written or generated? I believe FftOp length must be 1-d:
https://github.com/GleasonK/stablehlo/blob/main/docs/spec.md#inputs-34

%0 = "stablehlo.fft"(%input) {
    fft_length = dense<32> : tensor<1xi64>, fft_type = #stablehlo<fft_type RFFT>
} : (tensor<32xf32>) -> tensor<17xcomplex<f32>>

[GleasonK]

Regradless shouldn't crash.. Likely something missing in the FftOp verifier. Will dig into this.

[kuhar]

It's based on an older MHLO test. Whatever the input is, it should error out gracefully instead of crashing, no?

[GleasonK]

This attribute now uses DenseI64ArrayAttr, so invalid representations of fft_length will error gracefully.