[SDXL] Failed to legalize operation 'torch.operator' that was explicitly marked illegal : torch.aten._safe_softmax

[TCHeish]

CMD:

python3 models/turbine_models/custom_models/sdxl_inference/vae.py --precision=fp32 --external_weights=safetensors --device=cpu --compile_to=vmfb --iree_target_triple=x86_64-linux-gnu --external_weight_path=./vae_decode.safetensors

ERROR:

Compiling to llvm-cpu with flags: ['--iree-llvmcpu-target-triple=x86_64-linux-gnu', '--iree-llvmcpu-target-cpu-features=host', '--iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false', '--iree-llvmcpu-distribution-size=32', '--iree-opt-const-eval=false', '--iree-llvmcpu-enable-ukernels=all', '--iree-global-opt-enable-quantized-matmul-reassociation']
Saved to stable_diffusion_xl_base_1_0_bs1_1024x1024_fp32_vae_decode_decomp_.mlir
Traceback (most recent call last):
  File "/SHARK-ModelDev/models/turbine_models/custom_models/sdxl_inference/vae.py", line 204, in <module>
    mod_str = export_vae_model(
              ^^^^^^^^^^^^^^^^^
  File "/SHARK-ModelDev/models/turbine_models/custom_models/sdxl_inference/vae.py", line 177, in export_vae_model
    vmfb_path = utils.compile_to_vmfb(
                ^^^^^^^^^^^^^^^^^^^^^^
  File "/SHARK-ModelDev/models/turbine_models/custom_models/sd_inference/utils.py", line 320, in compile_to_vmfb
    flatbuffer_blob = ireec.compile_str(
                      ^^^^^^^^^^^^^^^^^^
  File "/SHARK-ModelDev/shark_venv_iree/lib/python3.11/site-packages/iree/compiler/tools/core.py", line 299, in compile_str
    result = invoke_immediate(cl, immediate_input=input_bytes)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/SHARK-ModelDev/shark_venv_iree/lib/python3.11/site-packages/iree/compiler/tools/binaries.py", line 198, in invoke_immediate
    raise CompilerToolError(process)
iree.compiler.tools.binaries.CompilerToolError: Error invoking IREE compiler tool iree-compile
Error code: 1
Diagnostics:
<stdin>:494:12: error: failed to legalize operation 'torch.operator' that was explicitly marked illegal
    %144 = torch.operator "torch.aten._safe_softmax"(%143, %int-1_169, %none_170) : (!torch.vtensor<[1,1,16384,16384],f32>, !torch.int, !torch.none) -> !torch.vtensor<[1,1,16384,16384],f32> 
           ^

Invoked with:
 iree-compile /SHARK-ModelDev/shark_venv_iree/lib/python3.11/site-packages/iree/compiler/tools/../_mlir_libs/iree-compile - --iree-input-type=torch --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-embedded-linker-path=/scratch/darrenh/work/shark/SHARK-ModelDev/shark_venv_iree/lib/python3.11/site-packages/iree/compiler/tools/../_mlir_libs/iree-lld --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-triple=x86_64-linux-gnu --iree-llvmcpu-target-cpu-features=host --iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false --iree-llvmcpu-distribution-size=32 --iree-opt-const-eval=false --iree-llvmcpu-enable-ukernels=all --iree-global-opt-enable-quantized-matmul-reassociation

Need more information? Set IREE_SAVE_TEMPS=/some/dir in your environment to save all artifacts and reproducers.

[TCHeish]

fallback torch from 2.5.0 to 2.3.0 to workaround the issue.