[RFC] Proposal for converting arith to linalg generic pass

MLIR
1

Background

The arithmetic dialect supports point-wise operations on tensors. However, arithmetic operations with tensor operands are not bufferizable. If the MLIR byte code contains arithmetic operations with tensor operands during the one-shot-bufferization pass then the one-shot-bufferization pass will fail.

Proposal

We would like to contribute a pass that replaces arith ops on tensors with arith ops on scalars embedded in the linalg::GenericOp. If this pass is run before the on-shot-bufferization pass then the semantics of the arith ops on tensors are lowered to linalg ops which allow for the bufferization of tensor operands.

Prototype

We’ve worked on a prototype that implements a functionality similar to that of TosaToLinalg conversion, however the prototype differs on some key aspects:

  • Instead of translating the TOSA dialect, the prototype translates only a subset of the arith dialect to linalg.
  • Whenever dynamic dimensions are encountered in the arith operands a comparison and an assertion are emitted to assert at run time that the dynamic dimensions are equal.
  • Operations between tensors must be of the same rank.

The current prototype could be refactored to share common code with TosaToLinalg but some direction is needed to place common infrastructure.

One can see the prototype here: Add -convert-arith-to-linalg pass. · erick-xanadu/llvm-project@8df9813 · GitHub

2

Hi - not sure if this fits the bill but have you looked at the lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp pass? I believe that it converts the strict subset of things that we believed we knew how to handle at the time that the ElementwiseMappable trait was defined.

3

Thanks!

I think this is exactly what I was hoping to implement. I’m wondering though, what are the semantics for tensors with dynamic dimensions when dimensions might differ? When I translate an example using -convert-elementwise-to-linalg:

I go from:

  func.func @transformaddfdynamic(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x4xf32> {
    %0 = arith.addf %arg0, %arg1 : tensor<?x4xf32>
    return %0 : tensor<?x4xf32>
  }

To:

  func.func @transformaddfdynamic(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x4xf32> {
    %0 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<?x4xf32>, tensor<?x4xf32>) outs(%arg0 : tensor<?x4xf32>) {
    ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
      %1 = arith.addf %arg2, %arg3 : f32
      linalg.yield %1 : f32
    } -> tensor<?x4xf32>
    return %0 : tensor<?x4xf32>
  }

which doesn’t show what will happen if the dimensions differ. Is this bit of semantic left to the next lowering passes?

4

Afk, but I believe the semantics of ElementwiseMappable are that different dimensions is UB (same as linalg). I remember there being a pass (or maybe just a discussion of a pass) which would add shape asserts at the linalg level. Most frontends are adding those at the top level, though (since in all generality, that is where you can reason about such things).

1 Like
5

In hlo land, we have broadcasting operations that support operands with different shapes (to some degree) and those are modeled in chlo, for example the BroadcastAddOp. When we lower these to mhlo which requires operands to have matching shape, shape constraints are inserted.

You could do something similar for the arith dialect but according to the documentation for the ElementwiseMappable trait (and as @stellaraccident mentioned), broadcasting is not allowed or supported. arith really just supports elementwise in the most elemental meaning and everything else is undefined.

1 Like