One-shot-bufferize pass fails due to the `builtin.unrealized_conversion_cast`

MLIR TOSA
1

one-shot-bufferize pass fails due to the builtin.unrealized_conversion_cast. By looking at the source changes, the below commit seems to be introducing this builtin.unrealized_conversion_cast.

Please look at changes to the TosaToLinalgPass, specifically the tosa::populateTosaTypeConversion change. populateTosaTypeConversion change is introduced as part of the below change.

tagging @mgehre-amd

Reproduction steps:

Input IR:

func.func @ui8_to_f32_dynamic(%arg0: tensor<1x?x?x3xui8>) -> (tensor<1x?x?x3xf32>)
 {
    %0 = tosa.cast %arg0 : (tensor<1x?x?x3xui8>) -> tensor<1x?x?x3xf32>
    return %0 : tensor<1x?x?x3xf32>
 }

The mlir-opt option to lower the program:

mlir-opt --pass-pipeline="builtin.module(func.func(tosa-to-linalg), one-shot-bufferize{allow-unknown-ops=true bufferize-function-boundaries=true}, func.func(finalizing-bufferize))" cast.tosa.mlir

Output IR for the passing case:

// -----// IR Dump After TosaToLinalg (tosa-to-linalg) //----- //
func.func @ui8_to_f32_dynamic(%arg0: tensor<1x?x?x3xui8>) -> tensor<1x?x?x3xf32> {
  %c1 = arith.constant 1 : index
  %dim = tensor.dim %arg0, %c1 : tensor<1x?x?x3xui8>
  %c2 = arith.constant 2 : index
  %dim_0 = tensor.dim %arg0, %c2 : tensor<1x?x?x3xui8>
  %0 = tensor.empty(%dim, %dim_0) : tensor<1x?x?x3xf32>
  %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<1x?x?x3xui8>) outs(%0 : tensor<1x?x?x3xf32>) {
  ^bb0(%in: ui8, %out: f32):
    %2 = builtin.unrealized_conversion_cast %in : ui8 to i8
    %3 = arith.uitofp %2 : i8 to f32
    linalg.yield %3 : f32
  } -> tensor<1x?x?x3xf32>
  return %1 : tensor<1x?x?x3xf32>
}

// -----// IR Dump After OneShotBufferize (one-shot-bufferize) //----- //
#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
  func.func @ui8_to_f32_dynamic(%arg0: memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) -> memref<1x?x?x3xf32> {
    %c1 = arith.constant 1 : index
    %dim = memref.dim %arg0, %c1 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
    %c2 = arith.constant 2 : index
    %dim_0 = memref.dim %arg0, %c2 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
    %alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<1x?x?x3xf32>
    linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) outs(%alloc : memref<1x?x?x3xf32>) {
    ^bb0(%in: ui8, %out: f32):
      %0 = builtin.unrealized_conversion_cast %in : ui8 to i8
      %1 = arith.uitofp %0 : i8 to f32
      linalg.yield %1 : f32
    }
    %cast = memref.cast %alloc : memref<1x?x?x3xf32> to memref<1x?x?x3xf32, strided<[?, ?, ?, ?], offset: ?>>
    return %alloc : memref<1x?x?x3xf32>
  }
}


// -----// IR Dump After FinalizingBufferize (finalizing-bufferize) //----- //
func.func @ui8_to_f32_dynamic(%arg0: memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) -> memref<1x?x?x3xf32> {
  %c1 = arith.constant 1 : index
  %dim = memref.dim %arg0, %c1 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
  %c2 = arith.constant 2 : index
  %dim_0 = memref.dim %arg0, %c2 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
  %alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<1x?x?x3xf32>
  linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) outs(%alloc : memref<1x?x?x3xf32>) {
  ^bb0(%in: ui8, %out: f32):
    %0 = builtin.unrealized_conversion_cast %in : ui8 to i8
    %1 = arith.uitofp %0 : i8 to f32
    linalg.yield %1 : f32
  }
  %cast = memref.cast %alloc : memref<1x?x?x3xf32> to memref<1x?x?x3xf32, strided<[?, ?, ?, ?], offset: ?>>
  return %alloc : memref<1x?x?x3xf32>
}

#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
  func.func @ui8_to_f32_dynamic(%arg0: memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) -> memref<1x?x?x3xf32> {
    %c1 = arith.constant 1 : index
    %dim = memref.dim %arg0, %c1 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
    %c2 = arith.constant 2 : index
    %dim_0 = memref.dim %arg0, %c2 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
    %alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<1x?x?x3xf32>
    linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) outs(%alloc : memref<1x?x?x3xf32>) {
    ^bb0(%in: ui8, %out: f32):
      %0 = builtin.unrealized_conversion_cast %in : ui8 to i8
      %1 = arith.uitofp %0 : i8 to f32
      linalg.yield %1 : f32
    }
    %cast = memref.cast %alloc : memref<1x?x?x3xf32> to memref<1x?x?x3xf32, strided<[?, ?, ?, ?], offset: ?>>
    return %alloc : memref<1x?x?x3xf32>
  }
}

Output IR for the failing case:

// -----// IR Dump After TosaToLinalg (tosa-to-linalg) //----- //
func.func @ui8_to_f32_dynamic(%arg0: tensor<1x?x?x3xui8>) -> tensor<1x?x?x3xf32> {
  %0 = builtin.unrealized_conversion_cast %arg0 : tensor<1x?x?x3xui8> to tensor<1x?x?x3xi8>
  %c1 = arith.constant 1 : index
  %dim = tensor.dim %0, %c1 : tensor<1x?x?x3xi8>
  %c2 = arith.constant 2 : index
  %dim_0 = tensor.dim %0, %c2 : tensor<1x?x?x3xi8>
  %1 = tensor.empty(%dim, %dim_0) : tensor<1x?x?x3xf32>
  %2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%0 : tensor<1x?x?x3xi8>) outs(%1 : tensor<1x?x?x3xf32>) {
  ^bb0(%in: i8, %out: f32):
    %3 = builtin.unrealized_conversion_cast %in : i8 to i8
    %4 = arith.uitofp %3 : i8 to f32
    linalg.yield %4 : f32
  } -> tensor<1x?x?x3xf32>
  return %2 : tensor<1x?x?x3xf32>
}

// -----// IR Dump After OneShotBufferize (one-shot-bufferize) //----- //
#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
  func.func @ui8_to_f32_dynamic(%arg0: memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) -> memref<1x?x?x3xf32> {
    %0 = bufferization.to_tensor %arg0 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
    %1 = builtin.unrealized_conversion_cast %0 : tensor<1x?x?x3xui8> to tensor<1x?x?x3xi8>
    %2 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
    %3 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
    %4 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
    %c1 = arith.constant 1 : index
    %dim = memref.dim %4, %c1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
    %c2 = arith.constant 2 : index
    %dim_0 = memref.dim %3, %c2 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
    %alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<1x?x?x3xf32>
    linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%2 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>) outs(%alloc : memref<1x?x?x3xf32>) {
    ^bb0(%in: i8, %out: f32):
      %5 = builtin.unrealized_conversion_cast %in : i8 to i8
      %6 = arith.uitofp %5 : i8 to f32
      linalg.yield %6 : f32
    }
    %cast = memref.cast %alloc : memref<1x?x?x3xf32> to memref<1x?x?x3xf32, strided<[?, ?, ?, ?], offset: ?>>
    return %alloc : memref<1x?x?x3xf32>
  }
}


cast.tosa.mlir:3:10: error: failed to legalize operation 'builtin.unrealized_conversion_cast'
    %0 = tosa.cast %arg0 : (tensor<1x?x?x3xui8>) -> tensor<1x?x?x3xf32>
         ^
cast.tosa.mlir:3:10: note: see current operation: %1 = "builtin.unrealized_conversion_cast"(%0) : (tensor<1x?x?x3xui8>) -> tensor<1x?x?x3xi8>
// -----// IR Dump After FinalizingBufferize Failed (finalizing-bufferize) //----- //
func.func @ui8_to_f32_dynamic(%arg0: memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>) -> memref<1x?x?x3xf32> {
  %0 = bufferization.to_tensor %arg0 : memref<1x?x?x3xui8, strided<[?, ?, ?, ?], offset: ?>>
  %1 = builtin.unrealized_conversion_cast %0 : tensor<1x?x?x3xui8> to tensor<1x?x?x3xi8>
  %2 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
  %3 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
  %4 = bufferization.to_memref %1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
  %c1 = arith.constant 1 : index
  %dim = memref.dim %4, %c1 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
  %c2 = arith.constant 2 : index
  %dim_0 = memref.dim %3, %c2 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>
  %alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<1x?x?x3xf32>
  linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%2 : memref<1x?x?x3xi8, strided<[?, ?, ?, ?], offset: ?>>) outs(%alloc : memref<1x?x?x3xf32>) {
  ^bb0(%in: i8, %out: f32):
    %5 = builtin.unrealized_conversion_cast %in : i8 to i8
    %6 = arith.uitofp %5 : i8 to f32
    linalg.yield %6 : f32
  }
  %cast = memref.cast %alloc : memref<1x?x?x3xf32> to memref<1x?x?x3xf32, strided<[?, ?, ?, ?], offset: ?>>
  return %alloc : memref<1x?x?x3xf32>
}
2

This appears related to Lowering tosa to linalg problem with unsigned types - #12 by mgehre-amd @mgehre-amd ?

3

I tried a bunch of Tosa operators with unsigned integer input, and all of them failed with the same issue. However, with the commented code below from the TosaTypeConverters, the problem goes away.

void mlir::tosa::populateTosaTypeConversion(TypeConverter &converter) {
  
  converter.addConversion([&](Type type) -> std::optional<Type> {
    if (type.isUnsignedInteger()) {
      return IntegerType::get(type.getContext(),
      type.getIntOrFloatBitWidth(),
                              IntegerType::SignednessSemantics::Signless);
    }
    return type;
  });
  // converter.addConversion([&](TensorType type) -> std::optional<Type> {
  //   auto converted = converter.convertType(type.getElementType());
  //   if (!converted)
  //     return {};
  //   return type.clone(converted);
  // });
  // converter.addSourceMaterialization([&](OpBuilder &builder, Type resultType,
  //                                        ValueRange inputs,
  //                                        Location loc) -> std::optional<Value> {
  //   if (inputs.size() != 1)
  //     return std::nullopt;

  //   return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
  //       .getResult(0);
  // });
  // converter.addTargetMaterialization([&](OpBuilder &builder, Type resultType,
  //                                        ValueRange inputs,
  //                                        Location loc) -> std::optional<Value> {
  //   if (inputs.size() != 1)
  //     return std::nullopt;
    
  //   return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
  //       .getResult(0);
  // });
}
4

Just in case someone has a similar issue, we can work around the problem by having a module-lowering pass like @MaheshRavishankar suggested in the other thread (Lowering tosa to linalg problem with unsigned types - #5 by MaheshRavishankar). Coincidentally, my colleague @rafaelubal also had the same idea. Essentially, we can get around this issue by converting function argument types from signed to signless representation and inserting the builtin.unrealized_conversion_cast operations in appropriate positions. And later, run the canonicalize or reconcile-unrealized-casts pass after the Tosa-to-linalg pass so that the redundant builtin.unrealized_conversion_cast operations get canceled. I will consider upstreaming this fix.