[all-commits] [llvm/llvm-project] 5f5b3b: [mlir][ArmSME] Add rewrites to swap extract of ext...

[Cullen Rhodes via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 5f5b3bb22b2e4ffcd14a8fc8a5edc14bc098a47e
      https://github.com/llvm/llvm-project/commit/5f5b3bb22b2e4ffcd14a8fc8a5edc14bc098a47e
  Author: Cullen Rhodes <cullen.rhodes at arm.com>
  Date:   2024-02-05 (Mon, 05 Feb 2024)

  Changed paths:
    M mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp
    M mlir/test/Dialect/ArmSME/outer-product-fusion.mlir

  Log Message:
  -----------
  [mlir][ArmSME] Add rewrites to swap extract of extend (#80407)

In mixed matmul lowering (e.g., i8 to i32) we're seeing the following
sequence:

  %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
  %1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
%lhs = vector.scalable.extract %1[0] : vector<[4]xi32> from
vector<[8]xi32>

  ... (same for rhs)

%2 = vector.outerproduct %lhs, %rhs, %acc vector<[4]xi32>,
vector<[4]xi32>

  // x4 chained by accumulator

This chain of 4 outer products can be fused into a single 4-way widening
variant but the pass doesn't match on the IR, as it expects the source
of the inputs to be an extend and it can't look through the extracts.

This patch fixes this with two rewrites that swaps extract(extend) into
extend(extract).

Related to #78975, #79288.