[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: add missing argument materialization. (PR #121200)

[Benjamin Chetioui]

bchetioui wrote:

> What kind of failure are you seeing? Is there no matching source materialization callback? It's possible that you have to turn an `addArgumentMaterialization` into an `addSourceMaterialization`.

This is the error I'm seeing:

```
error: failed to legalize unresolved materialization from ('f32') to ('tensor<f32>') that remained live after conversion
```

with IR

```
// -----// IR Dump After HloLegalizeToLinalgPass Failed (hlo-legalize-to-linalg) //----- //
func.func @map_mixed(%arg0: tensor<?xf32>, %arg1: tensor<4xf32>) -> tensor<?xf32> {
  %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<1xindex>
  %c0 = arith.constant 0 : index
  %extracted = tensor.extract %0[%c0] : tensor<1xindex>
  %1 = tensor.empty(%extracted) : tensor<?xf32>
  %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<?xf32>, tensor<4xf32>) outs(%1 : tensor<?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %3 = builtin.unrealized_conversion_cast %in_0 : f32 to tensor<f32>
    %4 = builtin.unrealized_conversion_cast %in : f32 to tensor<f32>
    %extracted_1 = tensor.extract %4[] : tensor<f32>
    %extracted_2 = tensor.extract %3[] : tensor<f32>
    %5 = arith.addf %extracted_1, %extracted_2 : f32
    %from_elements = tensor.from_elements %5 : tensor<f32>
    %extracted_3 = tensor.extract %from_elements[] : tensor<f32>
    linalg.yield %extracted_3 : f32
  } -> tensor<?xf32>
  return %2 : tensor<?xf32>
}
```

I believe that the materialization is inserted [here](https://github.com/llvm/llvm-project/blob/aa2fdc69d35ff1c4a6de8a8d8edcc4c15236bb15/mlir/lib/Transforms/Utils/DialectConversion.cpp#L1479-L1484). But admittedly, I'm not that familiar with these internals---let me try to look into `addArgumentMaterialization`.

https://github.com/llvm/llvm-project/pull/121200