llvm-project

9628061e - [mlir][AMDGPU] Add canonicalization pattern to pack scales for ScaledMFMAOp (#155951)

[mlir][AMDGPU] Add canonicalization pattern to pack scales for ScaledMFMAOp (#155951) The ScaledMFMAOp accepts scales as a vector of 4 bytes (`vector<4xf8E8M0FNU>`) that can be stored in a single register with a particular scale accessed using the `OpSel` attribute. Currently, we only use one byte in this 4-byte vector, resulting in 3 wasted registers. This is fixed by identifying when single byte extractions are performed and rewriting them into extractions of 4-byte vectors. Example: ``` %unit = vector.extract %ScaleSrc[offsets] : f8E8M0FNU from vector<?x?x?xf8E8M0FNU> %scale = vector.insert %unit, ... : f8E8M0FNU into vector<4xf8E8M0FNU> amdgpu.scaled_mfma(%scale[0] * ... ``` to ``` %reshaped = vector.shape_cast %ScaleSrc : vector<?x?x?xf8E8M0FNU> to vector<?x4xf8E8M0FNU> %scale = vector.extract %reshaped[?] : vector<4xf8E8M0FNU> from vector<?x4xf8E8M0FNU> amdgpu.scaled_mfma(%scale[0-3] * ... ``` --------- Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>