[compiler] SDXL EulerDiscreteScheduler Compilation Failure

[saienduri]

What happened?

When compiling the scheduler IR, we are seeing the following error:

/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:47:10: error: failed to legalize unresolved materialization from ('tensor<20xi32>') to ('tensor<?xi32>') that remained live after conversion
    %7 = torch.aten.argmax %6, %none, %false : !torch.vtensor<[20],si32>, !torch.none, !torch.bool -> !torch.vtensor<[],si64>
         ^
/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:47:10: note: see current operation: %41 = "builtin.unrealized_conversion_cast"(%40) : (tensor<20xi32>) -> tensor<?xi32>
/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:47:10: note: see existing live user here:
%26:2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%21 : tensor<?xi32>) outs(%25, %23 : tensor<i32>, tensor<i64>) {
^bb0(%in: i32, %out: i32, %out_39: i64):
  %52 = linalg.index 0 : index
  %53 = arith.index_cast %52 : index to i64
  %54 = arith.maxsi %in, %out : i32
  %55 = arith.cmpi sgt, %in, %out : i32
  %56 = arith.select %55, %53, %out_39 : i64
  linalg.yield %54, %56 : i32, i64
} -> (tensor<i32>, tensor<i64>)
/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:102:10: error: failed to legalize unresolved materialization from ('tensor<20xi32>') to ('tensor<?xi32>') that remained live after conversion
    %7 = torch.aten.argmax %6, %none, %false : !torch.vtensor<[20],si32>, !torch.none, !torch.bool -> !torch.vtensor<[],si64>
         ^
/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:102:10: note: see current operation: %52 = "builtin.unrealized_conversion_cast"(%51) : (tensor<20xi32>) -> tensor<?xi32>
/data/e2eshark/actions-runner-turbine/_work/SHARK-ModelDev/SHARK-ModelDev/vmfbs/stable_diffusion_xl_base_1_0_EulerDiscreteScheduler_bs1_1024x1024_fp16_20.mlir:102:10: note: see existing live user here:
%31:2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%26 : tensor<?xi32>) outs(%30, %28 : tensor<i32>, tensor<i64>) {
^bb0(%in: i32, %out: i32, %out_266: i64):
  %119 = linalg.index 0 : index
  %120 = arith.index_cast %119 : index to i64
  %121 = arith.maxsi %in, %out : i32
  %122 = arith.cmpi sgt, %in, %out : i32
  %123 = arith.select %122, %120, %out_266 : i64
  linalg.yield %121, %123 : i32, i64
} -> (tensor<i32>, tensor<i64>)

Steps to reproduce your issue

1. Please wget the following MLIR: https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sdxl-scheduler/11-21-2024/model.mlir

2. Run the following compile command:

iree-compile model.mlir --iree-hip-target=gfx942 --iree-hal-target-backends=rocm -o test.vmfb

What component(s) does this issue relate to?

Compiler

Version information

No response

Additional context

No response

[nithinsubbiah]

Related to https://github.com/iree-org/iree/pull/19171