Unverified Commit 5f5b3bb2 authored by Cullen Rhodes's avatar Cullen Rhodes Committed by GitHub
Browse files

[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.
parent 8e00fc33
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment