Brgemm register tiling support for bf16 type (#1005)

This PR extends the `brgemm register tiling` pass to support `bf16` type. The changes: 1) Template the existing pass to execute on `linalg.batch_reduce_matmul` for `fp32` and `linal.generic` for `vnni` opt bf16, 2) Test-cases for `bf16` type.
libxsmm · Feb 19, 2025 · f8d8a16 · f8d8a16
1 parent cb1e22f
commit f8d8a16
Show file tree

Hide file tree

Showing 7 changed files with 351 additions and 333 deletions.
diff --git a/benchmarks/config/base/base.json b/benchmarks/config/base/base.json
@@ -40,21 +40,21 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": ["avx512.*"]
     },
     "gemm_fp32_mlir_vector_avx2": {
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=4,16 '" ],
+      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=4,16,1 '" ],
       "extensions": ["avx2"]
     },
     "gemm_fp32_mlir_vector_sve": {
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=4,32 '" ],
+      "flags": [ "-n", "100",  "-run-args='--vector-to-kernels --registerBlocking=4,32,1 '" ],
       "extensions": ["asimd"]
     },
     "gemm_bf16_dp2_mlir": {
@@ -82,21 +82,21 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --bias --relu --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": ["avx512.*"]
     },
     "mlp_fp32_mlir_vector_avx2": {
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --bias --relu --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=4,16 '" ],
+      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=4,16,1 '" ],
       "extensions": ["avx2" ]
     },
     "mlp_fp32_mlir_vector_sve": {
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --bias --relu --float-type=f32 --batch=256 --layers=1024,1024,1024,1024 --tiles=32,32,32" ],
       "environment": {},
-      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=4,32 '" ],
+      "flags": [ "-n", "100",  "-run-args='--def-parallel --vector-to-kernels --registerBlocking=4,32,1 '" ],
       "extensions": ["asimd"]
     },
     "mlp_bf16_dp2_mlir": {
@@ -127,7 +127,7 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --float-type=f32 --batch=256 --layers=1024,1024,1024,1024" ],
       "environment": {},
-      "flags": [ "-n", "100", "-run-args='--vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100", "-run-args='--vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": [ "avx512.*" ]
     },
     "fp32_3x1024_args_mlir": {
@@ -141,7 +141,7 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=args --float-type=f32 --batch=256 --layers=1024,1024,1024,1024" ],
       "environment": {},
-      "flags": [ "-n", "100", "-run-args='--vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100", "-run-args='--vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": [ "avx512.*" ]
     },
     "bf16_3x1024_const_mlir": {
@@ -172,7 +172,7 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=const --bias --relu --float-type=f32 --batch=256 --layers=1024,1024,1024,1024" ],
       "environment": {},
-      "flags": [ "-n", "100", "-run-args='--def-parallel  --vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100", "-run-args='--def-parallel  --vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": [ "avx512.*" ]
     },
     "fp32_3x1024_args_mlir": {
@@ -186,7 +186,7 @@
       "type": "IR-GEN",
       "benchmark": [ "mlir-gen", "--kernel=args --bias --relu --float-type=f32 --batch=256 --layers=1024,1024,1024,1024" ],
       "environment": {},
-      "flags": [ "-n", "100", "-run-args=' --def-parallel  --vector-to-kernels --registerBlocking=8,32 '" ],
+      "flags": [ "-n", "100", "-run-args=' --def-parallel  --vector-to-kernels --registerBlocking=8,32,1 '" ],
       "extensions": [ "avx512.*" ]
     },
     "bf16_3x1024_const_mlir": {