[mlir][spirv] Error when using -convert-scf-to-spirv with message Failure to legalize unresolved materialization

[sweead]

Hello, I ran into an error when using mlir-opt -convert-scf-to-spirv on the following mixed IR.
To be specific, IR was successfully lowered to the cf dialect via -convert-scf-to-cf.
I'm not sure if this is a bug, it seems the problem is with type conversions.
IR:

//mlir-opt test.mlir -convert-scf-to-spirv

#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)>
module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spv.resource_limits<>>} {
  func.func @test_argmax(%arg0: tensor<14x19xf32>) {
    %c0_i32 = arith.constant 0 : i32
    %cst = arith.constant -3.40282347E+38 : f32
    %0 = bufferization.to_memref %arg0 : memref<14x19xf32>
    %1 = memref.alloc() {alignment = 128 : i64} : memref<14xi32>
    linalg.fill ins(%c0_i32 : i32) outs(%1 : memref<14xi32>)
    %2 = memref.alloc() {alignment = 128 : i64} : memref<14xf32>
    linalg.fill ins(%cst : f32) outs(%2 : memref<14xf32>)
    %3 = memref.alloc() {alignment = 128 : i64} : memref<14xi32>
    memref.copy %1, %3 : memref<14xi32> to memref<14xi32>
    %4 = memref.alloc() {alignment = 128 : i64} : memref<14xf32>
    memref.copy %2, %4 : memref<14xf32> to memref<14xf32>
    linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "reduction"]} ins(%0 : memref<14x19xf32>) outs(%3, %4 : memref<14xi32>, memref<14xf32>) {
    ^bb0(%arg1: f32, %arg2: i32, %arg3: f32):
      %5 = linalg.index 1 : index
      %6 = arith.index_cast %5 : index to i32
      %7 = arith.cmpf ogt, %arg1, %arg3 : f32
      %8 = arith.select %7, %arg1, %arg3 : f32
      %9 = arith.select %7, %6, %arg2 : i32
      linalg.yield %9, %8 : i32, f32
    }
    return
  }
}

error message:

test.mlir:19:12: error: failed to legalize unresolved materialization from 'index' to 'i32' that remained live after conversion
      %6 = arith.index_cast %5 : index to i32
           ^
test.mlir:19:12: note: see current operation: %10 = "builtin.unrealized_conversion_cast"(%9) : (index) -> i32
test.mlir:22:12: note: see existing live user here: %16 = "spv.Select"(%12, %10, %arg2) : (i1, i32, i32) -> i32
      %9 = arith.select %7, %6, %arg2 : i32
           ^

[llvmbot]

@llvm/issue-subscribers-mlir

[llvmbot]

@llvm/issue-subscribers-mlir-spirv

[kuhar]

Hi @sweead,

Thanks for reporting the issue. The IR used the old SPIR-V dialect prefix, so I upgraded it to something that can be processed by mlir today:

#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)>
module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>} {
  func.func @test_argmax(%arg0: tensor<14x19xf32>) {
    %c0_i32 = arith.constant 0 : i32
    %cst = arith.constant -3.40282347E+38 : f32
    %0 = bufferization.to_memref %arg0 : memref<14x19xf32>
    %1 = memref.alloc() {alignment = 128 : i64} : memref<14xi32>
    linalg.fill ins(%c0_i32 : i32) outs(%1 : memref<14xi32>)
    %2 = memref.alloc() {alignment = 128 : i64} : memref<14xf32>
    linalg.fill ins(%cst : f32) outs(%2 : memref<14xf32>)
    %3 = memref.alloc() {alignment = 128 : i64} : memref<14xi32>
    memref.copy %1, %3 : memref<14xi32> to memref<14xi32>
    %4 = memref.alloc() {alignment = 128 : i64} : memref<14xf32>
    memref.copy %2, %4 : memref<14xf32> to memref<14xf32>
    linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "reduction"]} ins(%0 : memref<14x19xf32>) outs(%3, %4 : memref<14xi32>, memref<14xf32>) {
    ^bb0(%arg1: f32, %arg2: i32, %arg3: f32):
      %5 = linalg.index 1 : index
      %6 = arith.index_cast %5 : index to i32
      %7 = arith.cmpf ogt, %arg1, %arg3 : f32
      %8 = arith.select %7, %arg1, %arg3 : f32
      %9 = arith.select %7, %6, %arg2 : i32
      linalg.yield %9, %8 : i32, f32
    }
    return
  }
}

I can reproduce the error, taking a look...

[kuhar]

@antiagainst I think it's another example of us doing conversion on ops within a loop without checking that the surrounding loop construct is supported (currently linalg.generic is not). I wonder if it'd be easier if we had an allow list passed as a predicate to tall of these terminal conversions.

[kuhar]

We have an existing TODO in the surrounding code:

  // TODO: Change SPIR-V conversion to be progressive and remove the following
  // patterns.
  mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns);
  populateFuncToSPIRVPatterns(typeConverter, patterns);
  populateMemRefToSPIRVPatterns(typeConverter, patterns);
  populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);

I guess this would entail introducing some level of nesting in how we run these helper patterns, e.g., arith.? Do we have some existing issue related to this TODO, @antiagainst?

[kuhar]

@sweead

I ran into an error when using mlir-opt -convert-scf-to-spirv on the following mixed IR.
To be specific, IR was successfully lowered to the cf dialect via -convert-scf-to-cf.

Yeah, like you noticed, SPIR-V currently expects some higher-level conversions to happen first. If this shows up in some real-world scenario you have to make sure to break down higher-level ops using other dialect conversions before running --convert-xyz-to-spirv passes.

[kuhar]

With the current conversion pass structure, this is working as intended -- we report a conversion error and don't crash. We have an easier spirv conversion pass pipeline in mind that should cover similar usecases.