'tensor.dim' op unexpected during shape cleanup; dynamic dimensions must have been resolved prior to leaving the flow dialect

[pdhirajkumarprasad]

for the given IR

module {
  func.func @"torch-jit-export"(%arg0: !torch.vtensor<[?,6],f32>) -> (!torch.vtensor<[?,1],si64> ) attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 21 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "1.7"} {
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<0> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64> 
    %2 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<4> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64> 
    %3 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<1> : tensor<1xsi64>} : () -> !torch.vtensor<[1],si64> 
    %4 = torch.operator "onnx.Slice"(%arg0, %1, %2, %3) : (!torch.vtensor<[?,6],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>, !torch.vtensor<[1],si64>) -> !torch.vtensor<[?,4],f32> 
    %5 = torch.operator "onnx.Softmax"(%4) {torch.onnx.axis = -1 : si64} : (!torch.vtensor<[?,4],f32>) -> !torch.vtensor<[?,4],f32> 
    %6 = torch.operator "onnx.Log"(%5) : (!torch.vtensor<[?,4],f32>) -> !torch.vtensor<[?,4],f32> 
    %7 = torch.operator "onnx.Multinomial"(%6) {torch.onnx.dtype = 7 : si64, torch.onnx.sample_size = 1 : si64} : (!torch.vtensor<[?,4],f32>) -> !torch.vtensor<[?,1],si64> 
    return %7 : !torch.vtensor<[?,1],si64>
  }
}

getting error as

model.torch_onnx.mlir:9:10: error: 'tensor.dim' op unexpected during shape cleanup; dynamic dimensions must have been resolved prior to leaving the flow dialect
    %7 = torch.operator "onnx.Multinomial"(%6) {torch.onnx.dtype = 7 : si64, torch.onnx.sample_size = 1 : si64} : (!torch.vtensor<[?,4],f32>) -> !torch.vtensor<[?,1],si64> 
         ^

command : iree-compile --iree-hal-target-backends=llvm-cpu -o abc.vmfb model.torch_onnx.mlir --iree-llvmcpu-target-cpu=host

[zjgarvey]

Here is a rather small reproducer. This is a simple cumulative sum performed as an scf loop. It fails CleanupTensorShapesPass during iree-compile. The tensor.extract and tensor.insert for the output tensor (the loop iter_arg) creates a tensor.dim op when converting tensor to flow.

module {
  func.func @forward(%arg0: tensor<?xf64>) -> tensor<?xf64> {
    %c0_i64 = arith.constant 0 : i64
    %c1_i64 = arith.constant 1 : i64
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?xf64>
    %0 = arith.index_cast %dim : index to i64
    %1 = tensor.empty(%dim) : tensor<?xf64>
    %2 = scf.for %arg2 = %c0_i64 to %0 step %c1_i64 iter_args(%arg3 = %1) -> (tensor<?xf64>)  : i64 {
      %3 = arith.cmpi sgt, %arg2, %c0_i64 : i64
      %4 = arith.index_cast %arg2 : i64 to index
      %extracted = tensor.extract %arg0[%4] : tensor<?xf64>
      %5 = arith.subi %arg2, %c1_i64 : i64
      %6 = arith.maxsi %5, %c0_i64 : i64
      %7 = arith.index_cast %6 : i64 to index
      %extracted_0 = tensor.extract %arg3[%7] : tensor<?xf64>
      %8 = arith.addf %extracted, %extracted_0 : f64
      %9 = arith.select %3, %8, %extracted : f64
      %inserted = tensor.insert %9 into %arg3[%4] : tensor<?xf64>
      scf.yield %inserted : tensor<?xf64>
    }
    return %2 : tensor<?xf64>
  }
}