On Improving Arm SME Lowering Resilience in MLIR

MLIR
2

Hi Hugo,

Thank you for creating this issue - there’s a lot of really “good” problems to discuss :slight_smile: Also, it’s great to see that people are finding the ArmSME work helpful :pray:

There’s a few separate issues here. Could I suggest that we limit this one to “canonicalization” (your first issue) and open separate threads for the other things? There’s a lot to discuss :sweat_smile:

Let me reply to your first question.

To be perfectly honest, we very rarely contribute to the cse/canonicalize patterns and often work around them. The problem with those patterns is that they are introduced at a certain point in time with a certain context in mind. Put differently, how do we decide that a pattern is a canonicalization? One can argue that given your experience, that pattern shouldn’t be a canonicalization.

For reference, you are most likely using this (note “canonicalization” at the end of the TD sequence):

func.func @matmul(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) -> tensor<?x?xf32> {
  %res = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
                       outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
  return %res : tensor<?x?xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%module : !transform.any_op {transform.consumed}) {
    %matmul = transform.structured.match ops{["linalg.matmul"]} in %module
      : (!transform.any_op) -> !transform.any_op

    // Step 1: Tile for size [4] x [4], which corresponds to SVLs x SVLs, where
    // SVLs is the number of 32-bit elements in a vector of SVL bits.
    %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul[[4], [4], 1]
      : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)

    // Step 2: Vectorize.
    transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1]
      : !transform.any_op

    // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which
    // currently only supports memrefs.
    %bufferize = transform.bufferization.one_shot_bufferize %module
      {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op

    %func = transform.structured.match ops{["func.func"]} in %bufferize
      : (!transform.any_op) -> !transform.any_op

    // THIS IS KEY!!!!
    transform.apply_registered_pass "canonicalize" to %func : (!transform.any_op) -> !transform.any_op

    transform.yield
  }
}

And this is the output:

$ mlir-opt --transform-interpreter file.mlir
(...)
          %7 = vector.mask %6 { vector.transfer_read %subview_5[%c0, %c0], %cst {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<[4]x[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>
          %8 = arith.mulf %3, %5 : vector<[4]x[4]x1xf32>
          %9 = vector.create_mask %0, %1 : vector<[4]x[4]xi1>
          %10 = vector.shape_cast %8 : vector<[4]x[4]x1xf32> to vector<[4]x[4]xf32>
          %11 = arith.addf %7, %10 : vector<[4]x[4]xf32>
          %12 = arith.select %9, %11, %10 : vector<[4]x[4]xi1>, vector<[4]x[4]xf32>
          vector.mask %6 { vector.transfer_write %12, %subview_5[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>> } : vector<[4]x[4]xi1>
(...)

Now, going back to your question:

IMHO, no. That would mean “lifting” abstractions. Instead, we should leverage the concept of “progressive lowering”. In this case, it feels like “canonicalization” has done too much. @MacDue has kindly dug out the PR that introduce this:

@vmurali + @dcaballe - what are your thoughts on this? Should we re-consider whether those patterns qualify as canonicalizations?

@nujaa , are you OK to open more threads for the other items to discuss?

-Andrzej