[KB][x86] AMX lowering pipelines

examples/KernelBench/level1.yaml

@@ -1,36 +1,36 @@
 - kernel: level1/1_Square_matrix_multiplication_.py
-  input_shapes: [1024x1024, 1024x1024]
+  input_shapes: [4096x4096, 4096x4096]
   initializations: [rnd, id]
-  output_shape: 1024x1024
-  gflops: (1024 * 1024 * 1024 * 2) / 1e9
+  output_shape: 4096x4096
+  gflops: (4096 * 4096 * 4096 * 2) / 1e9
   pipeline: matmul
 
 - kernel: level1/2_Standard_matrix_multiplication_.py
-  input_shapes: [512x1024, 1024x512]
+  input_shapes: [2048x8192, 8192x4096]
   initializations: [rnd, rnd]
-  output_shape: 512x512
-  gflops: (512 * 1024 * 512 * 2) / 1e9
+  output_shape: 2048x4096
+  gflops: (2048 * 8192 * 4096 * 2) / 1e9
   pipeline: matmul
 
 - kernel: level1/3_Batched_matrix_multiplication.py
-  input_shapes: [8x64x32, 8x32x64]
+  input_shapes: [3x2048x4096, 3x4096x2048]
   initializations: [rnd, rnd]
-  output_shape: 8x64x64
-  gflops: (8 * 64 * 32 * 64 * 2) / 1e9
+  output_shape: 3x2048x2048
+  gflops: (3 * 2048 * 4096 * 2048 * 2) / 1e9
   pipeline: batch_matmul
 
 - kernel: level1/4_Matrix_vector_multiplication_.py
-  input_shapes: [1024x1024, 1024x1]
+  input_shapes: [2048x8192, 8192x1]
   initializations: [rnd, rnd]
-  output_shape: 1024x1
-  gflops: (1024 * 1024 * 1 * 2) / 1e9
+  output_shape: 2048x1
+  gflops: (2048 * 8192 * 1 * 2) / 1e9
   pipeline: matvec
 
 - kernel: level1/5_Matrix_scalar_multiplication.py
-  input_shapes: [1024x1024, "1"]
+  input_shapes: [1024x16384, "1"]
   initializations: [rnd, rnd]
-  output_shape: 1024x1024
-  gflops: (1024 * 1024) / 1e9
+  output_shape: 1024x16384
+  gflops: (1024 * 16384) / 1e9
   pipeline: element_wise
 
 - kernel: level1/6_Matmul_with_large_K_dimension_.py
@@ -41,10 +41,10 @@
   pipeline: matmul
 
 - kernel: level1/7_Matmul_with_small_K_dimension_.py
-  input_shapes: [1024x64, 64x1024]
+  input_shapes: [4096x64, 64x4096]
   initializations: [rnd, rnd]
-  output_shape: 1024x1024
-  gflops: (1024 * 64 * 1024 * 2) / 1e9
+  output_shape: 4096x4096
+  gflops: (4096 * 64 * 4096 * 2) / 1e9
   pipeline: matmul
 
 - kernel: level1/8_Matmul_with_irregular_shapes_.py
@@ -55,17 +55,17 @@
   pipeline: matmul
 
 - kernel: level1/9_Tall_skinny_matrix_multiplication_.py
-  input_shapes: [1024x32, 32x1024]
+  input_shapes: [4096x64, 64x4096]
   initializations: [rnd, rnd]
-  output_shape: 1024x1024
-  gflops: (1024 * 32 * 1024 * 2) / 1e9
-  pipeline: batch_matmul
+  output_shape: 4096x4096
+  gflops: (4096 * 64 * 4096 * 2) / 1e9
+  pipeline: matmul
 
 - kernel: level1/10_3D_tensor_matrix_multiplication.py
-  input_shapes: [16x1024x2048, 2048x768]
+  input_shapes: [16x4096x2048, 2048x768]
   initializations: [rnd, rnd]
-  output_shape: 16x1024x768
-  gflops: (16 * 1024 * 2048 * 768 * 2) / 1e9
+  output_shape: 16x4096x768
+  gflops: (16 * 4096 * 2048 * 768 * 2) / 1e9
   pipeline: batch_matmul
 
 - kernel: level1/11_4D_tensor_matrix_multiplication.py
@@ -123,89 +123,89 @@
   pipeline: matmul
 
 - kernel: level1/19_ReLU.py
-  input_shapes: [1024x1024]
+  input_shapes: [4096x8192]
   initializations: [rnd]
-  output_shape: 1024x1024
-  gflops: (1024 * 1024) / 1e9
+  output_shape: 4096x8192
+  gflops: (1 * 4096 * 8192) / 1e9
 
 - kernel: level1/20_LeakyReLU.py
-  input_shapes: [1024x1024]
+  input_shapes: [4096x8192]
   initializations: [rnd]
-  output_shape: 1024x1024
-  gflops: (1024 * 1024) / 1e9
+  output_shape: 4096x8192
+  gflops: (2 * 4096 * 8192) / 1e9
   pipeline: element_wise
 
 - kernel: level1/21_Sigmoid.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (4 * 4096 * 8192) / 1e9
 
 - kernel: level1/22_Tanh.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (2 * 4096 * 8192) / 1e9
 
 - kernel: level1/23_Softmax.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (5 * 4096 * 8192) / 1e9
 
 - kernel: level1/24_LogSoftmax.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (5 * 4096 * 8192) / 1e9
 
 - kernel: level1/25_Swish.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (5 * 4096 * 8192) / 1e9
 
 - kernel: level1/26_GELU_.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (8 * 4096 * 8192) / 1e9
 
 - kernel: level1/27_SELU_.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (2 * 4096 * 8192) / 1e9
 
 - kernel: level1/28_HardSigmoid.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (2 * 4096 * 8192) / 1e9
 
 - kernel: level1/29_Softplus.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (4 * 4096 * 8192) / 1e9
 
 - kernel: level1/30_Softsign.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (2 * 4096 * 8192) / 1e9
 
 - kernel: level1/31_ELU.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (2 * 4096 * 8192) / 1e9
 
 - kernel: level1/32_HardTanh.py
   input_shapes: [4096x8192]
   initializations: [rnd]
   output_shape: 4096x8192
-  gflops: (4096 * 8192) / 1e9
+  gflops: (1 * 4096 * 8192) / 1e9
 
 - kernel: level1/33_BatchNorm.py
   input_shapes: [64x64x32x32]

examples/KernelBench/level3.yaml

@@ -0,0 +1,20 @@
+- kernel: level3/1_MLP.py
+  input_shapes: [128x16384]
+  initializations: [rnd]
+  output_shape: 128x8192
+  gflops: (2 * (128 * 16384 * 16384) + (128 * 16384 * 8192) + (128 * 8192)) / 1e9
+  pipeline: matmul
+
+- kernel: level3/2_ShallowWideMLP.py
+  input_shapes: [128x16384]
+  initializations: [rnd]
+  output_shape: 128x16384
+  gflops: (2 * (128 * 16384 * 16384) + (128 * 16384 * 16384) + (128 * 16384)) / 1e9
+  pipeline: matmul
+
+- kernel: level3/3_DeepNarrowMLP.py
+  input_shapes: [1024x8192]
+  initializations: [rnd]
+  output_shape: 1024x8192
+  gflops: (17 * (1024 * 1024 * 8192) + (1024 * 8192)) / 1e9
+  pipeline: matmul

examples/KernelBench/schedules/x86_64/amx/batch_matmul/bf16.yaml

@@ -0,0 +1,18 @@
+# This is an optimizing pipeline for simple GEMM kernels.
+Pipeline:
+  - include: ../../matmul-like.yaml {
+                                block_factors=[32,32,32]
+                                tile_size=32
+                                reg_tile_batch=1
+                                reg_tile_m=32
+                                reg_tile_n=32
+                                reg_tile_k=32
+                                batch=1
+                                reg_unroll_m=16
+                                reg_unroll_n=16
+                                reg_unroll_k=32
+                                fill_tiling=[1]
+                                generic_tiling=[1,1,1,1,8]
+                                2D_generic_tiling=[1,8]
+                                leading_batch_dims=1
+                              }

examples/KernelBench/schedules/x86_64/element_wise/bf16.yaml → examples/KernelBench/schedules/x86_64/amx/element_wise/bf16.yaml

@@ -1,6 +1,6 @@
 Pipeline:
   # Element-wise is usually done at F32 precision anyway
-  - include: common.yaml{
+  - include: ../../element-wise.yaml {
                           register_tile=16
                           unroll_factor=16
                         }

examples/KernelBench/schedules/x86_64/amx/matmul/bf16.yaml

@@ -0,0 +1,18 @@
+# This is an optimizing pipeline for simple GEMM kernels.
+Pipeline:
+  - include: ../../matmul-like.yaml {
+                                block_factors=[32,32,32]
+                                tile_size=32
+                                reg_tile_batch=1
+                                reg_tile_m=32
+                                reg_tile_n=32
+                                reg_tile_k=32
+                                batch=1
+                                reg_unroll_m=16
+                                reg_unroll_n=16
+                                reg_unroll_k=32
+                                fill_tiling=[1]
+                                generic_tiling=[1,1,1,8]
+                                2D_generic_tiling=[1,8]
+                                leading_batch_dims=0
+                              }

examples/KernelBench/schedules/x86_64/matvec/bf16.yaml → examples/KernelBench/schedules/x86_64/amx/matvec/bf16.yaml

@@ -1,5 +1,5 @@
 Pipeline:
-  - include: common.yaml{
+  - include: ../../mat-vec.yaml {
                           cache_tile=64
                           register_tile=8
                           unroll_factor=4

examples/KernelBench/schedules/x86_64/batch_matmul/f32.yaml

@@ -1,14 +1,18 @@
 # This is an optimizing pipeline for simple GEMM kernels.
 Pipeline:
   - include: ../matmul-like.yaml {
-                                reg_tile_batch=0
+                                block_factors=[32,32,32]
+                                tile_size=32
+                                reg_tile_batch=1
                                 reg_tile_m=8
                                 reg_tile_n=32
                                 reg_tile_k=2
-                                batch=0
+                                batch=1
                                 reg_unroll_m=1
                                 reg_unroll_n=16
                                 reg_unroll_k=1
-                                fill_tiling=[1,1,1]
-                                generic_tiling=[1,8]
+                                fill_tiling=[1]
+                                generic_tiling=[1,1,1,1,8]
+                                2D_generic_tiling=[1,8]
+                                leading_batch_dims=1
                               }

examples/KernelBench/schedules/x86_64/element_wise/common.yaml → examples/KernelBench/schedules/x86_64/element-wise.yaml

@@ -13,4 +13,4 @@ Pipeline:
 
   - include: bufferization.yaml
 
-  - include: ../lower.yaml
+  - include: lower.yaml

examples/KernelBench/schedules/x86_64/element_wise/f32.yaml

@@ -1,5 +1,5 @@
 Pipeline:
-  - include: common.yaml{
+  - include: ../element-wise.yaml {
                           register_tile=16
                           unroll_factor=16
                         }

examples/KernelBench/schedules/x86_64/lower.yaml

@@ -5,6 +5,7 @@ Pipeline:
   - include: bufferization_cleanup.yaml
 
   ## Buffer vectorization (gets rid of all the rest)
+  - schedule: "x86/amx_move_offsets.py[gen=amx_move_offsets]"
   - schedule: "vectorization.py[gen=x86_vectorization]"
   - schedule: "vectorization.py[gen=vectorize_all]"
   - schedule: "vectorization.py[gen=flatten_vector_ops]"

examples/KernelBench/schedules/x86_64/matvec/common.yaml → examples/KernelBench/schedules/x86_64/mat-vec.yaml

@@ -30,4 +30,4 @@ Pipeline:
   - pass: "scf-parallel-loop-fusion"
   - pass: "scf-parallel-loop-tiling"
 
-  - include: ../lower.yaml
+  - include: lower.yaml

examples/KernelBench/schedules/x86_64/matmul-like.yaml

@@ -1,6 +1,10 @@
 # This is an optimizing pipeline for kernel_bench matmul-like kernels.
 Pipeline:
-  - include: pack_and_tile.yaml
+  - include: pack_and_tile.yaml {
+                                  block_factors=$block_factors
+                                  tile_size=$tile_size
+                                  leading_batch_dims=$leading_batch_dims
+                                }
 
   ## CPU specific register tiling (depends on uArch & data type)
   - schedule: "x86/register_tiling.py[gen=matmul_register_tiling] {
@@ -20,8 +24,10 @@ Pipeline:
                                                                     reg_unroll_n=$reg_unroll_n
                                                                     reg_unroll_k=$reg_unroll_k
                                                                   }"
+  - pass: "linalg-morph-ops{named-to-category generic-to-named}"
   - schedule: "tiling.py[gen=tile_ops]{target_op=linalg.fill tile_sizes=$fill_tiling filter_ops=True}"
   - schedule: "tiling.py[gen=tile_ops]{target_op=linalg.generic tile_sizes=$generic_tiling filter_ops=True}"
+  - schedule: "tiling.py[gen=tile_ops]{target_op=linalg.generic tile_sizes=$2D_generic_tiling filter_ops=True}"
 
   - include: vectorize.yaml
 

examples/KernelBench/schedules/x86_64/matmul/f32.yaml

@@ -1,6 +1,8 @@
 # This is an optimizing pipeline for simple GEMM kernels.
 Pipeline:
   - include: ../matmul-like.yaml {
+                                block_factors=[32,32,32]
+                                tile_size=32
                                 reg_tile_batch=1
                                 reg_tile_m=8
                                 reg_tile_n=32
@@ -9,6 +11,8 @@ Pipeline:
                                 reg_unroll_m=1
                                 reg_unroll_n=16
                                 reg_unroll_k=1
-                                fill_tiling=[1,1,1]
-                                generic_tiling=[1,8]
+                                fill_tiling=[1]
+                                generic_tiling=[1,1,1,8]
+                                2D_generic_tiling=[1,8]
+                                leading_batch_dims=0
                               }