[transform] Better default tiles for small parallel dims

lighthouse/dialects/transform/transform_ext/utils/tiling/common.py

@@ -15,6 +15,16 @@ def parallel_and_reduction_dims(out_map: ir.AffineMap) -> tuple[list[int], list[
     return parallel_dims, reduction_dims
 
 
+def largest_divisor(dim_size: int | None, limit: int) -> int:
+    """Largest divisor of `dim_size` not above `limit` (1 for unknown sizes)."""
+    if dim_size is None:
+        return 1
+    for tile in range(min(dim_size, max(1, limit)), 0, -1):
+        if dim_size % tile == 0:
+            return tile
+    return 1
+
+
 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:

lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py

@@ -4,7 +4,7 @@
 from lighthouse.execution.target import TargetInfo
 from lighthouse.utils.mlir import linalg_inputs, linalg_outputs, opview
 
-from .common import parallel_and_reduction_dims
+from .common import largest_divisor, parallel_and_reduction_dims
 
 
 def _contraction_operand_types(
@@ -66,6 +66,38 @@ def vector_lane_count(target: TargetInfo | None, elem_type: ir.Type) -> int:
     return 16
 
 
+def _vector_tiles(dim_sizes: list[int], vector_width: int) -> list[int]:
+    """Tiles of the trailing dims spanning one vector register.
+
+    Inner dims narrower than the vector are taken whole and the next dims
+    tiled by a divisor of the lanes left, so a narrow innermost dim does not
+    shrink the tiles to a few elements.
+    """
+    if dim_sizes:
+        innermost_size = dim_sizes[-1]
+        if (
+            ir.ShapedType.is_dynamic_size(innermost_size)
+            or innermost_size >= vector_width
+        ):
+            return [vector_width]
+
+    tiles: list[int] = []
+    remaining = vector_width
+    for dim_size in reversed(dim_sizes):
+        dynamic = ir.ShapedType.is_dynamic_size(dim_size)
+        if dynamic:
+            tiles.append(1)
+            break
+        if dim_size >= remaining:
+            tiles.append(largest_divisor(dim_size, remaining))
+            break
+        tiles.append(max(1, dim_size))
+        remaining //= tiles[-1]
+        if remaining <= 1:
+            break
+    return tiles[::-1]
+
+
 def generic_parallel_tiles(
     op: ir.Operation | ir.OpView,
     out_map: ir.AffineMap,
@@ -75,9 +107,9 @@ def generic_parallel_tiles(
     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]
+    out_type = ir.ShapedType(linalg_outputs(op)[0].type)
+    vector_width = vector_lane_count(target, out_type.element_type)
+    return _vector_tiles(out_type.shape, vector_width)
 
 
 def generic_reduction_tiles() -> list[int]:

test/transform/test_tile_and_fuse_register_unroll.py

@@ -349,6 +349,59 @@ def assign_elementwise_register_unroll64():
 """
 
 
+ELTWISE_NARROW_INNER_DIMS = """
+#id = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#perm = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
+module {
+  func.func @even_inner(%x: tensor<1x1024x2xf32>) -> tensor<1x1024x2xbf16> {
+    %e = tensor.empty() : tensor<1x1024x2xbf16>
+    %a = linalg.generic {indexing_maps = [#id, #id], iterator_types = ["parallel", "parallel", "parallel"]}
+        ins(%x : tensor<1x1024x2xf32>) outs(%e : tensor<1x1024x2xbf16>) {
+    ^bb0(%i: f32, %o: bf16):
+      %v = arith.truncf %i : f32 to bf16
+      linalg.yield %v : bf16
+    } -> tensor<1x1024x2xbf16>
+    return %a : tensor<1x1024x2xbf16>
+  }
+
+  func.func @odd_inner(%x: tensor<8x40x3xf32>) -> tensor<8x40x3xf32> {
+    %e = tensor.empty() : tensor<8x40x3xf32>
+    %a = linalg.generic {indexing_maps = [#id, #id], iterator_types = ["parallel", "parallel", "parallel"]}
+        ins(%x : tensor<8x40x3xf32>) outs(%e : tensor<8x40x3xf32>) {
+    ^bb0(%i: f32, %o: f32):
+      %v = math.exp %i : f32
+      linalg.yield %v : f32
+    } -> tensor<8x40x3xf32>
+    return %a : tensor<8x40x3xf32>
+  }
+
+  func.func @dynamic_outer(%x: tensor<4x?x2xf32>) -> tensor<4x?x2xf32> {
+    %c1 = arith.constant 1 : index
+    %d = tensor.dim %x, %c1 : tensor<4x?x2xf32>
+    %e = tensor.empty(%d) : tensor<4x?x2xf32>
+    %a = linalg.generic {indexing_maps = [#id, #id], iterator_types = ["parallel", "parallel", "parallel"]}
+        ins(%x : tensor<4x?x2xf32>) outs(%e : tensor<4x?x2xf32>) {
+    ^bb0(%i: f32, %o: f32):
+      %v = math.exp %i : f32
+      linalg.yield %v : f32
+    } -> tensor<4x?x2xf32>
+    return %a : tensor<4x?x2xf32>
+  }
+
+  func.func @permuted_map(%x: tensor<32x2x4xi8>) -> tensor<32x2x4xi8> {
+    %e = tensor.empty() : tensor<32x2x4xi8>
+    %a = linalg.generic {indexing_maps = [#perm, #perm], iterator_types = ["parallel", "parallel", "parallel"]}
+        ins(%x : tensor<32x2x4xi8>) outs(%e : tensor<32x2x4xi8>) {
+    ^bb0(%i: i8, %o: i8):
+      %v = math.absi %i : i8
+      linalg.yield %v : i8
+    } -> tensor<32x2x4xi8>
+    return %a : tensor<32x2x4xi8>
+  }
+}
+"""
+
+
 # CHECK-LABEL: Test: tile_and_unroll_annotated
 # CHECK: func.func @main
 # CHECK-NOT: scf.for
@@ -598,3 +651,22 @@ def assign_elementwise_register_unroll64():
         ELTWISE_DTYPE_WIDTHS,
         assign_elementwise_register_unroll64,
     )
+
+
+# Inner dims narrower than a vector are taken whole and the next dim fills the
+# remaining lanes with a divisor of its dim size; a dynamic dim stays untiled.
+# CHECK-LABEL: Test: eltwise_narrow_inner_dims_register_unroll
+# CHECK: func.func @even_inner
+# CHECK: transform_ext.tile_sizes = array<i64: 1, 16, 2>
+# CHECK: func.func @odd_inner
+# CHECK: transform_ext.tile_sizes = array<i64: 1, 5, 3>
+# CHECK: func.func @dynamic_outer
+# CHECK: transform_ext.tile_sizes = array<i64: 1, 1, 2>
+# CHECK: func.func @permuted_map
+# CHECK: transform_ext.tile_sizes = array<i64: 2, 4, 8>
+with TargetInfo.override(features=["avx512f"]):
+    run(
+        "eltwise_narrow_inner_dims_register_unroll",
+        ELTWISE_NARROW_INNER_DIMS,
+        assign_elementwise_register_unroll,
+    )