diff --git a/examples/KernelBench/level1.yaml b/examples/KernelBench/level1.yaml index f7fb6e12..85d05cdf 100644 --- a/examples/KernelBench/level1.yaml +++ b/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] diff --git a/examples/KernelBench/level3.yaml b/examples/KernelBench/level3.yaml new file mode 100644 index 00000000..0bd0f360 --- /dev/null +++ b/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 diff --git a/examples/KernelBench/schedules/x86_64/amx/batch_matmul/bf16.yaml b/examples/KernelBench/schedules/x86_64/amx/batch_matmul/bf16.yaml new file mode 100644 index 00000000..98ff873b --- /dev/null +++ b/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 + } diff --git a/examples/KernelBench/schedules/x86_64/element_wise/bf16.yaml b/examples/KernelBench/schedules/x86_64/amx/element_wise/bf16.yaml similarity index 82% rename from examples/KernelBench/schedules/x86_64/element_wise/bf16.yaml rename to examples/KernelBench/schedules/x86_64/amx/element_wise/bf16.yaml index de9219c5..a668969e 100644 --- a/examples/KernelBench/schedules/x86_64/element_wise/bf16.yaml +++ b/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 } diff --git a/examples/KernelBench/schedules/x86_64/amx/matmul/bf16.yaml b/examples/KernelBench/schedules/x86_64/amx/matmul/bf16.yaml new file mode 100644 index 00000000..e912046c --- /dev/null +++ b/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 + } diff --git a/examples/KernelBench/schedules/x86_64/matvec/bf16.yaml b/examples/KernelBench/schedules/x86_64/amx/matvec/bf16.yaml similarity index 82% rename from examples/KernelBench/schedules/x86_64/matvec/bf16.yaml rename to examples/KernelBench/schedules/x86_64/amx/matvec/bf16.yaml index f44a9335..32e39b82 100644 --- a/examples/KernelBench/schedules/x86_64/matvec/bf16.yaml +++ b/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 diff --git a/examples/KernelBench/schedules/x86_64/batch_matmul/bf16.yaml b/examples/KernelBench/schedules/x86_64/batch_matmul/bf16.yaml deleted file mode 100644 index 8a25a63f..00000000 --- a/examples/KernelBench/schedules/x86_64/batch_matmul/bf16.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# This is an optimizing pipeline for simple GEMM kernels. -Pipeline: - - include: ../matmul-like.yaml { - reg_tile_batch=0 - reg_tile_m=8 - reg_tile_n=32 - reg_tile_k=2 - batch=0 - reg_unroll_m=1 - reg_unroll_n=16 - reg_unroll_k=2 - fill_tiling=[1,1,1] - generic_tiling=[1,8] - } diff --git a/examples/KernelBench/schedules/x86_64/batch_matmul/f32.yaml b/examples/KernelBench/schedules/x86_64/batch_matmul/f32.yaml index 1a778445..e22eb4d4 100644 --- a/examples/KernelBench/schedules/x86_64/batch_matmul/f32.yaml +++ b/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 } diff --git a/examples/KernelBench/schedules/x86_64/element_wise/common.yaml b/examples/KernelBench/schedules/x86_64/element-wise.yaml similarity index 96% rename from examples/KernelBench/schedules/x86_64/element_wise/common.yaml rename to examples/KernelBench/schedules/x86_64/element-wise.yaml index 331e85e2..ac255cde 100644 --- a/examples/KernelBench/schedules/x86_64/element_wise/common.yaml +++ b/examples/KernelBench/schedules/x86_64/element-wise.yaml @@ -13,4 +13,4 @@ Pipeline: - include: bufferization.yaml - - include: ../lower.yaml + - include: lower.yaml diff --git a/examples/KernelBench/schedules/x86_64/element_wise/f32.yaml b/examples/KernelBench/schedules/x86_64/element_wise/f32.yaml index 0655dc3f..836d9341 100644 --- a/examples/KernelBench/schedules/x86_64/element_wise/f32.yaml +++ b/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 } diff --git a/examples/KernelBench/schedules/x86_64/lower.yaml b/examples/KernelBench/schedules/x86_64/lower.yaml index 6dedcece..97f16772 100644 --- a/examples/KernelBench/schedules/x86_64/lower.yaml +++ b/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]" diff --git a/examples/KernelBench/schedules/x86_64/matvec/common.yaml b/examples/KernelBench/schedules/x86_64/mat-vec.yaml similarity index 98% rename from examples/KernelBench/schedules/x86_64/matvec/common.yaml rename to examples/KernelBench/schedules/x86_64/mat-vec.yaml index fd601a74..6d6fc9fa 100644 --- a/examples/KernelBench/schedules/x86_64/matvec/common.yaml +++ b/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 diff --git a/examples/KernelBench/schedules/x86_64/matmul-like.yaml b/examples/KernelBench/schedules/x86_64/matmul-like.yaml index d68e1e95..f8b37f7d 100644 --- a/examples/KernelBench/schedules/x86_64/matmul-like.yaml +++ b/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 diff --git a/examples/KernelBench/schedules/x86_64/matmul/bf16.yaml b/examples/KernelBench/schedules/x86_64/matmul/bf16.yaml deleted file mode 100644 index 24f96083..00000000 --- a/examples/KernelBench/schedules/x86_64/matmul/bf16.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# This is an optimizing pipeline for simple GEMM kernels. -Pipeline: - - include: ../matmul-like.yaml { - reg_tile_batch=1 - reg_tile_m=8 - reg_tile_n=32 - reg_tile_k=2 - batch=1 - reg_unroll_m=1 - reg_unroll_n=16 - reg_unroll_k=2 - fill_tiling=[1,1,1] - generic_tiling=[1,8] - } diff --git a/examples/KernelBench/schedules/x86_64/matmul/f32.yaml b/examples/KernelBench/schedules/x86_64/matmul/f32.yaml index 588d0d2e..4b7d5fff 100644 --- a/examples/KernelBench/schedules/x86_64/matmul/f32.yaml +++ b/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 } diff --git a/examples/KernelBench/schedules/x86_64/matvec/f32.yaml b/examples/KernelBench/schedules/x86_64/matvec/f32.yaml index 946f2527..03fa59df 100644 --- a/examples/KernelBench/schedules/x86_64/matvec/f32.yaml +++ b/examples/KernelBench/schedules/x86_64/matvec/f32.yaml @@ -1,5 +1,5 @@ Pipeline: - - include: common.yaml{ + - include: ../mat-vec.yaml { cache_tile=32 register_tile=8 unroll_factor=4 diff --git a/examples/KernelBench/schedules/x86_64/pack_and_tile.yaml b/examples/KernelBench/schedules/x86_64/pack_and_tile.yaml index 25d28da3..597961bf 100644 --- a/examples/KernelBench/schedules/x86_64/pack_and_tile.yaml +++ b/examples/KernelBench/schedules/x86_64/pack_and_tile.yaml @@ -3,8 +3,8 @@ Pipeline: ## Packing & Cache tiling (CPU generic) - pass: "linalg-fuse-elementwise-ops" - - schedule: "packing.py[gen=block_pack_matmuls]{block_factors=[32,32,32] rhs_transpose_outer_block=True rhs_transpose_inner_block=False}" - - schedule: "x86/pack_lowering.py[gen=lower_packs_unpacks]{tile_size=32}" + - schedule: "packing.py[gen=block_pack_matmuls]{block_factors=$block_factors rhs_transpose_outer_block=True rhs_transpose_inner_block=False}" + - schedule: "x86/pack_lowering.py[gen=lower_packs_unpacks]{tile_size=$tile_size leading_batch_dims=$leading_batch_dims}" - pass: "linalg-morph-ops{named-to-category generic-to-category}" - schedule: "x86/cache_tiling.py[gen=matmul_cache_tiling]{target=linalg.contract fuse_producers}" - schedule: "linalg.py[gen=linalg_contract_fold_unit_dims]" diff --git a/examples/KernelBench/schedules/x86_64/vectorize.yaml b/examples/KernelBench/schedules/x86_64/vectorize.yaml index 57abb67c..0a6609cb 100644 --- a/examples/KernelBench/schedules/x86_64/vectorize.yaml +++ b/examples/KernelBench/schedules/x86_64/vectorize.yaml @@ -4,4 +4,3 @@ Pipeline: - schedule: "vectorization.py[gen=vectorize_linalg]" - schedule: "hoisting.py[gen=hoist_loops]" - schedule: "vectorization.py[gen=simplify_vector_ops]" - - schedule: "vectorization.py[gen=x86_vectorization]" diff --git a/examples/KernelBench/test-kernel-bench.py b/examples/KernelBench/test-kernel-bench.py index 75b0ca5f..d8952d7c 100755 --- a/examples/KernelBench/test-kernel-bench.py +++ b/examples/KernelBench/test-kernel-bench.py @@ -21,6 +21,7 @@ yaml_files = [ script_path / "level1.yaml", script_path / "level2.yaml", + script_path / "level3.yaml", ] ci_files = [ script_path / "ci.yaml", @@ -39,7 +40,32 @@ def __init__(self, arch: str = None, feature: str = None): self.arch = arch else: self.arch = platform.machine() - self.feature = feature + + self.feature = [] + if feature is not None: + self.feature = feature.split(",") + else: + self._get_extension_list() + + def _get_extension_list(self) -> list[str]: + flags = subprocess.run( + "lscpu | grep Flags", + capture_output=True, + text=True, + shell=True, + ).stdout + schedule_path = script_path / "schedules" / self.arch + if not schedule_path.exists(): + return [] + + extensions = [] + for item in schedule_path.iterdir(): + if item.is_dir(): + extensions.append(item.name) + + for ext in extensions: + if ext in flags: + self.feature.append(ext) def get_pipeline_file(name: str, dtype: str, target: TargetInfo) -> Path: @@ -51,8 +77,10 @@ def get_pipeline_file(name: str, dtype: str, target: TargetInfo) -> Path: return kb_default_pipeline # First check if there's a pipeline file specific to the target features. - if target.feature: - pipeline = script_path / f"schedules/{target.feature}/{name}/{dtype}.yaml" + for feature in target.feature: + pipeline = ( + script_path / f"schedules/{target.arch}/{feature}/{name}/{dtype}.yaml" + ) if pipeline.exists(): return pipeline @@ -219,11 +247,14 @@ def get_flops_per_second(stdout: str, gflops: float) -> float: command_line += ["--target", target_info.arch] if target_info.feature: - command_line += ["--feature", target_info.feature] + command_line += ["--feature", ",".join(target_info.feature)] # Benchmark mode. if args.benchmark: command_line += ["--benchmark"] + # FIXME: This is here just for quick testing + # Remove when merging back to main + command_line += ["--nwarmup", "5", "--nruns", "10", "--no-validate"] # Shape inference or from args. if not args.infer_shapes: diff --git a/lighthouse/schedule/x86/pack_lowering.py b/lighthouse/schedule/x86/pack_lowering.py index d2f16c39..7d5e0bd3 100644 --- a/lighthouse/schedule/x86/pack_lowering.py +++ b/lighthouse/schedule/x86/pack_lowering.py @@ -56,9 +56,9 @@ def lower_unpacks_for_vectorization( tiled_unpack = structured.TileUsingForOp( unpack_op, sizes=unpack_tile_sizes ).tiled_linalg_op - if vector_tile_sizes: + for size in vector_tile_sizes: tiled_unpack = structured.TileUsingForOp( - tiled_unpack, sizes=vector_tile_sizes + tiled_unpack, sizes=[size] ).tiled_linalg_op structured.structured_lower_unpack( transform.OperationType.get("tensor.empty"), @@ -72,37 +72,51 @@ def lower_unpacks_for_vectorization( transform.yield_() -def lower_packs_unpacks(tile_size: int) -> ir.Module: +def lower_packs_unpacks(tile_size: int, leading_batch_dims: int = 0) -> ir.Module: """ Lower pack and unpack ops into hardware-friendly shapes. Args: tile_size: Target shape for sub-tiling pack and unpack ops' inner tiles + leading_batch_dims: Number of leading batch dimensions for extra tiling Returns: Schedule """ + batch_tile_sizes = [1] * leading_batch_dims with schedule_boilerplate() as (schedule, named_seq): pack_unpack_vector_m = max(8, tile_size) pack_unpack_vector_n = min(64, tile_size) packs = lh_transform.match_op(named_seq.bodyTarget, "linalg.pack") lower_packs_for_vectorization( pack_ops=packs, - pack_tile_sizes=[1, 1], - vector_tile_sizes=[1, 1, pack_unpack_vector_m, pack_unpack_vector_n], - vector_unroll_factors=[tile_size // pack_unpack_vector_n], + pack_tile_sizes=batch_tile_sizes + [1, 1], + vector_tile_sizes=batch_tile_sizes + + [1, 1, pack_unpack_vector_m, pack_unpack_vector_n], + vector_unroll_factors=batch_tile_sizes + + [tile_size // pack_unpack_vector_n], ) lh_transform.cleanup(named_seq.bodyTarget) unpacks = lh_transform.match_op(named_seq.bodyTarget, "linalg.unpack") lower_unpacks_for_vectorization( unpack_ops=unpacks, - unpack_tile_sizes=[tile_size, tile_size], - vector_tile_sizes=[1], + unpack_tile_sizes=batch_tile_sizes + [tile_size, tile_size], + vector_tile_sizes=batch_tile_sizes + [1], ) transposes = lh_transform.match_op(named_seq.bodyTarget, "linalg.transpose") with lh_transform.foreach(transposes) as tranpose: + tranpose = structured.TileUsingForOp( + tranpose, sizes=batch_tile_sizes + [1, 1, 1] + ).tiled_linalg_op structured.structured_vectorize(tranpose, []) transform.yield_() + copies = lh_transform.match_op(named_seq.bodyTarget, "linalg.copy") + with lh_transform.foreach(copies) as copy: + copy = structured.TileUsingForOp( + copy, sizes=batch_tile_sizes + [1] + ).tiled_linalg_op + structured.structured_vectorize(copy, []) + transform.yield_() # Cleanup. with ir.InsertionPoint( diff --git a/lighthouse/transform/vectorization.py b/lighthouse/transform/vectorization.py index 225c71a2..9e9b01c7 100644 --- a/lighthouse/transform/vectorization.py +++ b/lighthouse/transform/vectorization.py @@ -26,6 +26,7 @@ def x86_vector_patterns(target): target: Handle to target """ with ir.InsertionPoint(transform.ApplyPatternsOp(target).patterns): + x86.apply_patterns_x86_vector_contract_to_amx_dot_product() x86.apply_patterns_x86_vector_contract_to_packed_type_dot_product() x86.apply_patterns_x86_vector_contract_to_fma() x86.apply_patterns_x86_sink_vector_producer_ops()