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

[Markus Böck]

zero9178 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`.

Could you try adjusting this line: https://github.com/openxla/xla/blob/eb9d08bae564680ff465d772ceb70f4d84542e8c/xla/mlir_hlo/mhlo/utils/type_conversion.cc#L112
to also add the same function as a source materialization? 

As far as I can see in your IR, the only difference between before and now is that the `builtin.unrealized_conversion_cast` seen here is a source rather than argument materialization

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