[Mlir-commits] [mlir] [mlir][LinalgBlockPackMatmul] Add support for scalable block factors (PR #211354)

[Adam Siemieniuk]

================
@@ -0,0 +1,91 @@
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=[32],[16],[64] allow-padding=1" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=SCALABLE
+
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=[32],[16],[64] allow-padding=0" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=SCALABLE-NOPAD
+
+// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=32,16,[64] allow-padding=1" \
+// RUN: -canonicalize -split-input-file | FileCheck %s --check-prefix=MIXED
+
+// -----
+
+// All-scalable block factors, static input shapes, allow-padding=1.
+// Inner tile sizes and outer dimensions are all dynamic (vscale-relative).
+
+func.func @block_matmul_scalable_static(
+    %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {
+  %0 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>
+  return %0 : tensor<128x128xf32>
+}
+
+// SCALABLE-LABEL: func @block_matmul_scalable_static(
+// SCALABLE-DAG: %[[C32:.+]] = arith.constant 32 : index
+// SCALABLE-DAG: %[[C16:.+]] = arith.constant 16 : index
+// SCALABLE-DAG: %[[C64:.+]] = arith.constant 64 : index
+// SCALABLE-DAG: %[[MB:.+]] = arith.muli %{{.+}}, %[[C32]] : index
----------------
adam-smnk wrote:

I'd expect to see at least one `vscale` somewhere.
You could add `-cse` to make checks more readable as well.

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