From e05589dcf1df49dc9a93a054e1a0a64ff36e2671 Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Tue, 11 Aug 2026 11:44:30 +0200 Subject: [PATCH] [transform] Target-aware register parallel dim tiling strategy Adds register-level parallel strategy that derives tile sizes based on operations and TargetInfo. Includes shared target-capability analysis and tile-assignment helpers. Register-level strategy which directly mimics existing first-level register tiling which focuses on parallel dimension to ensure data can fit in available registers. The strategy folds in the two existing variants: general f32 and AMX bf16, and provides a generic target-aware SIMD size fallback. Assisted-by: Claude --- .../transform_ext/utils/tiling/common.py | 44 +++++ .../transform_ext/utils/tiling/registry.py | 2 + .../tiling/strategy_register_parallel.py | 43 +++++ .../transform_ext/utils/tiling/target_caps.py | 74 ++++++++ test/transform/test_assign_tile_sizes.py | 5 + .../test_tile_and_fuse_register_tiling.py | 147 ++++++++++++++++ test/transform/test_tile_size_ops.py | 15 ++ .../test_tiling_strategy_register_parallel.py | 50 ++++++ .../test_tiling_strategy_target_aware.py | 160 ++++++++++++++++++ 9 files changed, 540 insertions(+) create mode 100644 lighthouse/dialects/transform/transform_ext/utils/tiling/strategy_register_parallel.py create mode 100644 lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py create mode 100644 test/transform/test_tile_and_fuse_register_tiling.py create mode 100644 test/transform/test_tiling_strategy_register_parallel.py create mode 100644 test/transform/test_tiling_strategy_target_aware.py diff --git a/lighthouse/dialects/transform/transform_ext/utils/tiling/common.py b/lighthouse/dialects/transform/transform_ext/utils/tiling/common.py index 1d64d62b..4f5c126b 100644 --- a/lighthouse/dialects/transform/transform_ext/utils/tiling/common.py +++ b/lighthouse/dialects/transform/transform_ext/utils/tiling/common.py @@ -15,6 +15,50 @@ def parallel_and_reduction_dims(out_map: ir.AffineMap) -> tuple[list[int], list[ return parallel_dims, reduction_dims +def assign_from_tail(dims: list[int], values: list[int], sizes: list[int]) -> None: + """Write `values` onto `sizes` at the trailing `dims`, aligning both by their tails.""" + if not dims or not values: + return + tail = dims[-len(values) :] + vals = values[-len(tail) :] + for dim, tile in zip(tail, vals): + sizes[dim] = int(tile) + + +def assign_parallel_tiles( + parallel_dims: list[int], inner_tiles: list[int], sizes: list[int] +) -> None: + """Outer parallel dims -> 1; innermost parallel tail -> inner_tiles.""" + inner_count = min(len(inner_tiles), len(parallel_dims)) + for d in parallel_dims[:-inner_count] if inner_count else []: + sizes[d] = 1 + assign_from_tail(parallel_dims, inner_tiles, sizes) + + +def assign_reduction_tiles( + reduction_dims: list[int], red_tiles: list[int], sizes: list[int] +) -> None: + """Innermost reduction tail -> red_tiles; remaining reduction dims -> 1.""" + assign_from_tail(reduction_dims, red_tiles, sizes) + if reduction_dims: + protected = set(reduction_dims[-len(red_tiles) :]) if red_tiles else set() + unitize_unassigned_dims(reduction_dims, sizes, protected=protected) + + +def unitize_unassigned_dims( + dims: list[int], + sizes: list[int], + protected: set[int] | None = None, +) -> None: + """Set each untiled (zero) dim to 1, leaving `protected` dims untouched.""" + protected = protected or set() + for dim in dims: + if dim in protected: + continue + if sizes[dim] == 0: + sizes[dim] = 1 + + def output_tensor_dim_of_iter_dim(out_map: ir.AffineMap) -> dict[int, int]: """Map each iteration dim to the output tensor dim it indexes (if any).""" mapping: dict[int, int] = {} diff --git a/lighthouse/dialects/transform/transform_ext/utils/tiling/registry.py b/lighthouse/dialects/transform/transform_ext/utils/tiling/registry.py index 534eb1a7..18d60405 100644 --- a/lighthouse/dialects/transform/transform_ext/utils/tiling/registry.py +++ b/lighthouse/dialects/transform/transform_ext/utils/tiling/registry.py @@ -1,10 +1,12 @@ from .strategy_base import TilingStrategy from .strategy_cache import CacheTilingStrategy +from .strategy_register_parallel import RegisterParallelTilingStrategy # Maps each canonical strategy name to its implementation class. _STRATEGY_REGISTRY: dict[str, type[TilingStrategy]] = { "cache": CacheTilingStrategy, + "register_parallel": RegisterParallelTilingStrategy, } diff --git a/lighthouse/dialects/transform/transform_ext/utils/tiling/strategy_register_parallel.py b/lighthouse/dialects/transform/transform_ext/utils/tiling/strategy_register_parallel.py new file mode 100644 index 00000000..10f1a11d --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/tiling/strategy_register_parallel.py @@ -0,0 +1,43 @@ +from mlir import ir + +from lighthouse.utils.mlir import opview + +from .strategy_base import StrategyContext, TilingStrategy +from .common import ( + assign_parallel_tiles, + disable_small_tiles, + parallel_and_reduction_dims, +) +from .target_caps import ( + generic_parallel_tiles, + is_amx_bf16_contraction, + is_f32_contraction, +) + + +class RegisterParallelTilingStrategy(TilingStrategy): + """Register-level tiling of parallel dimensions; target-derived defaults.""" + + def compute( + self, op: ir.Operation | ir.OpView, ctx: StrategyContext + ) -> list[int] | None: + out_map = self.output_map(op) + if out_map is None: + return None + + sizes = [0] * out_map.n_dims + parallel_dims, _ = parallel_and_reduction_dims(out_map) + if not parallel_dims: + return None + + ov = opview(op) + if is_amx_bf16_contraction(ov, ctx.target): + inner_tiles = [32, 32] + elif is_f32_contraction(ov): + inner_tiles = [8, 32] + else: + inner_tiles = generic_parallel_tiles(ov, out_map, ctx.target) + + assign_parallel_tiles(parallel_dims, inner_tiles, sizes) + disable_small_tiles(ov, out_map, sizes, max(inner_tiles, default=ctx.tile_size)) + return sizes diff --git a/lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py b/lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py new file mode 100644 index 00000000..83b6e6a8 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py @@ -0,0 +1,74 @@ +from mlir import ir +from mlir.dialects import linalg + +from lighthouse.execution.target import TargetInfo +from lighthouse.utils.mlir import linalg_inputs, linalg_outputs, opview + +from .common import parallel_and_reduction_dims + + +def _contraction_operand_types( + op: ir.Operation | ir.OpView, +) -> tuple[ir.Type, ir.Type, ir.Type] | None: + """(lhs, rhs, acc) element types of a single-output contraction, else None.""" + ov = opview(op) + if not linalg.isa_contraction_op(ov): + return None + inputs = linalg_inputs(ov) + outputs = linalg_outputs(ov) + if inputs is None or outputs is None or len(inputs) < 2 or len(outputs) != 1: + return None + return ( + ir.ShapedType(inputs[0].type).element_type, + ir.ShapedType(inputs[1].type).element_type, + ir.ShapedType(outputs[0].type).element_type, + ) + + +def is_amx_bf16_contraction( + op: ir.Operation | ir.OpView, target: TargetInfo | None +) -> bool: + """True for a bf16 -> f32 contraction on an AMX-capable target.""" + if target is None or not target.is_supported("amx"): + return False + types = _contraction_operand_types(op) + if types is None: + return False + lhs, rhs, acc = types + return ( + isinstance(lhs, ir.BF16Type) + and isinstance(rhs, ir.BF16Type) + and isinstance(acc, ir.F32Type) + ) + + +def is_f32_contraction(op: ir.Operation | ir.OpView) -> bool: + """True for a contraction with all-f32 operands (lhs, rhs and acc).""" + types = _contraction_operand_types(op) + return types is not None and all(isinstance(t, ir.F32Type) for t in types) + + +def _vector_lane_count(target: TargetInfo | None, elem_type: ir.Type) -> int: + """SIMD lane count for `elem_type` on `target` (defaults assume 512-bit).""" + vector_bits = ( + target.vector_register_width_bits + if target is not None and target.vector_register_width_bits is not None + else 512 + ) + if isinstance(elem_type, (ir.FloatType, ir.IntegerType)): + return max(1, vector_bits // max(1, elem_type.width)) + return 16 + + +def generic_parallel_tiles( + op: ir.Operation | ir.OpView, + out_map: ir.AffineMap, + target: TargetInfo | None, +) -> list[int]: + """SIMD-lane default parallel tiles for ops without a microkernel profile.""" + parallel_dims, _ = parallel_and_reduction_dims(out_map) + if not parallel_dims: + return [] + out_elem = ir.ShapedType(linalg_outputs(op)[0].type).element_type + inner = _vector_lane_count(target, out_elem) + return [inner] if len(parallel_dims) == 1 else [1, inner] diff --git a/test/transform/test_assign_tile_sizes.py b/test/transform/test_assign_tile_sizes.py index 4da97b63..beb9b066 100644 --- a/test/transform/test_assign_tile_sizes.py +++ b/test/transform/test_assign_tile_sizes.py @@ -66,6 +66,11 @@ def run( print(payload) +# CHECK-LABEL: Test: strategy_attr_register_parallel +# CHECK: linalg.matmul +# CHECK-SAME: transform_ext.tile_sizes = array +run("strategy_attr_register_parallel", PAYLOAD, "linalg.matmul", "register_parallel") + # CHECK-LABEL: Test: strategy_attr_cache # CHECK: linalg.matmul # CHECK-SAME: transform_ext.tile_sizes = array diff --git a/test/transform/test_tile_and_fuse_register_tiling.py b/test/transform/test_tile_and_fuse_register_tiling.py new file mode 100644 index 00000000..cea6b0c9 --- /dev/null +++ b/test/transform/test_tile_and_fuse_register_tiling.py @@ -0,0 +1,147 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir import ir + +import lighthouse.dialects as lh_dialects +from lighthouse.execution.target import TargetInfo +from lighthouse.schedule import tile_and_fuse as tf + + +def run(name: str, payload_str: str, *schedules): + print(f"Test: {name}", flush=True) + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + payload = ir.Module.parse(payload_str) + modules = [] + for make_schedule in schedules: + sched = make_schedule() + modules.append(sched) + sched.body.operations[0].apply(payload.operation) + print(payload) + + +def run_with_target_override( + name: str, + payload_str: str, + *, + features: list[str] | None = None, + arch: str | None = None, + schedules, +): + with TargetInfo.override(features=features, arch=arch): + run(name, payload_str, *schedules) + + +def assign_register_parallel(): + return tf.assign_and_propagate_tile_sizes( + tile_size=32, + strategy="register_parallel", + ) + + +def tile_and_fuse_keep(): + return tf.tile_and_fuse_annotated(clear_annotations=False) + + +MLP = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +#map1 = affine_map<(d0, d1) -> (d1)> +module { + func.func @main(%arg0: tensor<32x64xf32>, %w: tensor<64x128xf32>, %b: tensor<128xf32>) + -> tensor<32x128xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %1 = tensor.empty() : tensor<32x128xf32> + %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<32x128xf32>) -> tensor<32x128xf32> + %3 = linalg.matmul ins(%arg0, %w : tensor<32x64xf32>, tensor<64x128xf32>) + outs(%2 : tensor<32x128xf32>) -> tensor<32x128xf32> + %4 = linalg.generic {indexing_maps = [#map, #map1, #map], + iterator_types = ["parallel", "parallel"]} + ins(%3, %b : tensor<32x128xf32>, tensor<128xf32>) + outs(%1 : tensor<32x128xf32>) { + ^bb0(%in: f32, %in_2: f32, %out: f32): + %6 = arith.addf %in, %in_2 : f32 + linalg.yield %6 : f32 + } -> tensor<32x128xf32> + %5 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%4 : tensor<32x128xf32>) + outs(%1 : tensor<32x128xf32>) { + ^bb0(%in: f32, %out: f32): + %6 = arith.cmpf ugt, %in, %cst : f32 + %7 = arith.select %6, %in, %cst : f32 + linalg.yield %7 : f32 + } -> tensor<32x128xf32> + return %5 : tensor<32x128xf32> + } +} +""" + +MLP_BF16 = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +#map1 = affine_map<(d0, d1) -> (d1)> +module { + func.func @main(%arg0: tensor<32x64xbf16>, %w: tensor<64x128xbf16>, %b: tensor<128xf32>) + -> tensor<32x128xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %1 = tensor.empty() : tensor<32x128xf32> + %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<32x128xf32>) -> tensor<32x128xf32> + %3 = linalg.matmul ins(%arg0, %w : tensor<32x64xbf16>, tensor<64x128xbf16>) + outs(%2 : tensor<32x128xf32>) -> tensor<32x128xf32> + %4 = linalg.generic {indexing_maps = [#map, #map1, #map], + iterator_types = ["parallel", "parallel"]} + ins(%3, %b : tensor<32x128xf32>, tensor<128xf32>) + outs(%1 : tensor<32x128xf32>) { + ^bb0(%in: f32, %in_2: f32, %out: f32): + %6 = arith.addf %in, %in_2 : f32 + linalg.yield %6 : f32 + } -> tensor<32x128xf32> + %5 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%4 : tensor<32x128xf32>) + outs(%1 : tensor<32x128xf32>) { + ^bb0(%in: f32, %out: f32): + %6 = arith.cmpf ugt, %in, %cst : f32 + %7 = arith.select %6, %in, %cst : f32 + linalg.yield %7 : f32 + } -> tensor<32x128xf32> + return %5 : tensor<32x128xf32> + } +} +""" + + +# CHECK-LABEL: Test: register_parallel_tile_and_fuse +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +# CHECK: scf.forall +run( + "register_parallel_tile_and_fuse", MLP, assign_register_parallel, tile_and_fuse_keep +) + + +# Register-parallel strategy can infer target defaults from features. +# CHECK-LABEL: Test: register_parallel_target_defaults_amx +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run_with_target_override( + "register_parallel_target_defaults_amx", + MLP_BF16, + features=["amx_tile"], + schedules=(assign_register_parallel, tile_and_fuse_keep), +) + + +# An AMX-capable target must not change f32 GEMM tiling: the AMX branch only +# applies to bf16 x bf16 -> f32 contractions, so f32 keeps its own defaults. +# CHECK-LABEL: Test: register_parallel_f32_under_amx_target +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run_with_target_override( + "register_parallel_f32_under_amx_target", + MLP, + features=["amx_tile"], + schedules=(assign_register_parallel, tile_and_fuse_keep), +) diff --git a/test/transform/test_tile_size_ops.py b/test/transform/test_tile_size_ops.py index 879cdf6e..2138fd28 100644 --- a/test/transform/test_tile_size_ops.py +++ b/test/transform/test_tile_size_ops.py @@ -193,3 +193,18 @@ def build_get_leading_unit_tile_sizes_schedule(op_name: str): BATCH_MATMUL_PAYLOAD, lambda: build_assign_schedule("linalg.batch_matmul"), ) + + +# Register-parallel strategy: batch outer dim tiled to 1, inner M/N tiled by +# f32 FMA defaults, K untiled. +# CHECK-LABEL: Test: register_parallel_strategy +# CHECK: linalg.batch_matmul +# CHECK-SAME: transform_ext.tile_sizes = array +run( + "register_parallel_strategy", + BATCH_MATMUL_PAYLOAD, + lambda: build_assign_schedule( + "linalg.batch_matmul", + strategy="register_parallel", + ), +) diff --git a/test/transform/test_tiling_strategy_register_parallel.py b/test/transform/test_tiling_strategy_register_parallel.py new file mode 100644 index 00000000..ac921d01 --- /dev/null +++ b/test/transform/test_tiling_strategy_register_parallel.py @@ -0,0 +1,50 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir import ir +from mlir.dialects import transform + +import lighthouse.dialects as lh_dialects +from lighthouse import transform as lh_transform +from lighthouse.dialects.transform.transform_ext import assign_tile_sizes +from lighthouse.schedule.builders import schedule_boilerplate + + +def run(name: str, payload_str: str, build_schedule): + print(f"Test: {name}", flush=True) + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + payload = ir.Module.parse(payload_str) + sched = build_schedule() + sched.body.operations[0].apply(payload.operation) + print(payload) + + +PAYLOAD = """ +module { + func.func @main(%a: tensor<4x64x64xf32>, %b: tensor<4x64x64xf32>) -> tensor<4x64x64xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<4x64x64xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<4x64x64xf32>) -> tensor<4x64x64xf32> + %mm = linalg.batch_matmul ins(%a, %b : tensor<4x64x64xf32>, tensor<4x64x64xf32>) + outs(%f : tensor<4x64x64xf32>) -> tensor<4x64x64xf32> + return %mm : tensor<4x64x64xf32> + } +} +""" + + +def build_schedule(): + with schedule_boilerplate() as (sched, named_seq): + ops = lh_transform.match_op(named_seq.bodyTarget, "linalg.batch_matmul") + assign_tile_sizes( + ops, + strategy="register_parallel", + ) + transform.yield_() + return sched + + +# CHECK-LABEL: Test: register_parallel_strategy +# CHECK: linalg.batch_matmul +# CHECK-SAME: transform_ext.tile_sizes = array +run("register_parallel_strategy", PAYLOAD, build_schedule) diff --git a/test/transform/test_tiling_strategy_target_aware.py b/test/transform/test_tiling_strategy_target_aware.py new file mode 100644 index 00000000..00331279 --- /dev/null +++ b/test/transform/test_tiling_strategy_target_aware.py @@ -0,0 +1,160 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir import ir +from mlir.dialects import transform + +import lighthouse.dialects as lh_dialects +from lighthouse import transform as lh_transform +from lighthouse.dialects.transform.transform_ext import assign_tile_sizes +from lighthouse.execution.target import TargetInfo +from lighthouse.schedule.builders import schedule_boilerplate + + +def run(name: str, payload_str: str, build_schedule): + print(f"Test: {name}", flush=True) + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + payload = ir.Module.parse(payload_str) + sched = build_schedule() + sched.body.operations[0].apply(payload.operation) + print(payload) + + +F32_MATMUL = """ +module { + func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + return %mm : tensor<128x128xf32> + } +} +""" + + +BF16_MATMUL = """ +module { + func.func @main(%a: tensor<128x64xbf16>, %b: tensor<64x128xbf16>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul ins(%a, %b : tensor<128x64xbf16>, tensor<64x128xbf16>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + return %mm : tensor<128x128xf32> + } +} +""" + + +# A non-contraction op, so register_parallel falls back to generic_parallel_tiles, +# whose inner tile is the target's SIMD lane count for the output element type. +ELTWISE = """ +module { + func.func @main(%a: tensor<64x64xf32>, %b: tensor<64x64xf32>) -> tensor<64x64xf32> { + %sum = linalg.add ins(%a, %b : tensor<64x64xf32>, tensor<64x64xf32>) + outs(%a : tensor<64x64xf32>) -> tensor<64x64xf32> + return %sum : tensor<64x64xf32> + } +} +""" + + +def build_register_parallel(): + with schedule_boilerplate() as (sched, named_seq): + ops = lh_transform.match_op(named_seq.bodyTarget, "linalg.matmul") + assign_tile_sizes( + ops, + strategy="register_parallel", + ) + transform.yield_() + return sched + + +# CHECK-LABEL: Test: f32_register_parallel_default +# CHECK: linalg.matmul +# CHECK-SAME: transform_ext.tile_sizes = array +run("f32_register_parallel_default", F32_MATMUL, lambda: build_register_parallel()) + + +# CHECK-LABEL: Test: bf16_amx_register_parallel_default +# CHECK: linalg.matmul +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=["amx_tile"]): + run( + "bf16_amx_register_parallel_default", + BF16_MATMUL, + lambda: build_register_parallel(), + ) + + +# Without AMX, bf16 matmul falls back to the generic SIMD-lane tiling. +# CHECK-LABEL: Test: bf16_no_amx_register_parallel_generic_fallback +# CHECK: linalg.matmul +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=[]): + run( + "bf16_no_amx_register_parallel_generic_fallback", + BF16_MATMUL, + lambda: build_register_parallel(), + ) + + +def build_register_parallel_eltwise(): + with schedule_boilerplate() as (sched, named_seq): + ops = lh_transform.match_op(named_seq.bodyTarget, "linalg.add") + assign_tile_sizes( + ops, + strategy="register_parallel", + ) + transform.yield_() + return sched + + +# 512-bit vectors (AVX-512): 32-bit lanes -> inner tile of 16. +# CHECK-LABEL: Test: eltwise_register_parallel_avx512 +# CHECK: linalg.add +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=["avx512f"]): + run( + "eltwise_register_parallel_avx512", + ELTWISE, + lambda: build_register_parallel_eltwise(), + ) + + +# 256-bit vectors (AVX2): 32-bit lanes -> inner tile of 8. +# CHECK-LABEL: Test: eltwise_register_parallel_avx2 +# CHECK: linalg.add +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=["avx2"]): + run( + "eltwise_register_parallel_avx2", + ELTWISE, + lambda: build_register_parallel_eltwise(), + ) + + +# 128-bit vectors (SSE): 32-bit lanes -> inner tile of 4. +# CHECK-LABEL: Test: eltwise_register_parallel_sse +# CHECK: linalg.add +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=["sse4_1"]): + run( + "eltwise_register_parallel_sse", + ELTWISE, + lambda: build_register_parallel_eltwise(), + ) + + +# No recognized vector extension: falls back to the 512-bit assumption -> 16. +# CHECK-LABEL: Test: eltwise_register_parallel_no_features +# CHECK: linalg.add +# CHECK-SAME: transform_ext.tile_sizes = array +with TargetInfo.override(features=[]): + run( + "eltwise_register_parallel_no_features", + ELTWISE, + lambda: build_register_parallel_eltwise(), + )