[Mlir-commits] [mlir] [mlir][ArmNeon] Add pack/mmt4d path to f32 linalg.matmul integration test (PR #225063)

[Andrzej Warzyński]

================
@@ -0,0 +1,316 @@
+// REQUIRES: target={{(aarch64|arm64).*}}
+
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE:   -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE:   -cse -canonicalize -convert-vector-to-scf \
+// DEFINE:   -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm \
+// DEFINE:   -o %t
+
+// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon" \
+// DEFINE:    -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+
+// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+/// HIGH-LEVEL OVERVIEW
+///
+/// End-to-end test for computing matrix-multiplication (f32, accumulating to
+/// f32) using linalg.mmt4d. In particular, demonstrates how the following
+/// MLIR sequence (implemented in @matmul_via_mmt4d):
+///
+///   A_pack = linalg.pack A
+///   B_pack = linalg.pack B
+///   C_pack = linalg.pack C
+///   out_pack = linalg.mmt4d(A_pack, B_pack, C_pack)
+///
+/// is equivalent to:
+///
+///  linalg.matmul(A, B, C)
+///
+/// (implemented in @matmul).
+///
+/// NOTES ON IMPLEMENTATION
+/// 1. @matmul is only lowered via `-test-lower-to-llvm`, with no tiling or
+///    vectorization -- it's there to provide a reference check, not to
+///    demonstrate a lowering path.
+///
+/// 2. NEON has no instruction that multiplies two f32 matrices in one go, so
+///    the inner tile is (M0, N0, K0) = (4, 4, 1): every vector.contract is a
+///    rank-1 update, lowered to outer products, i.e. one `fmla` per row of
+///    the 4x4 accumulator tile.
+///
+/// 3. Packing is still worthwhile without a dedicated instruction: it gives
+///    the tiles a statically-known shape (no masking in the mmt4d loop) and
+///    makes them contiguous in memory.
----------------
banach-space wrote:

[nit] I wouldn't worry about there being (or not) a dedicated matrix-multiplication instruction. As you correctly observe, `linalg.mmt4d` is helpful regardless (it enables packing and that alone is a worthwhile optimisation). 

Also, in this file we don't really bother checking whether relevant instructions are generated (well, since we rely on generic FMAs, there wouldn't be much to check). We do check things like that when it does become key/relevant, e.g. here: https://github.com/llvm/llvm-project/blob/387200eb1b8fc636730f819ba487e1191df6dff1/mlir/test/Integration/Dialect/Vector/CPU/ArmNeon/vector-contract-bfmmla.mlir?plain=1#L54

Basically, I would just skip this part and focus on documenting what is available and tested rather than what is not.

https://github.com/llvm/llvm-project/pull/225063