Add evoformer example

iree/turbine/kernel/_support/dtype.py

@@ -1,5 +1,6 @@
 __all__ = [
     "DataType",
+    "bf16",
     "bool",
     "i4",
     "i8",
@@ -17,7 +18,16 @@
 ]
 
 _INT_TYPES = ["i1", "i4", "i8", "i16", "i32", "i64"]
-_FLOAT_TYPES = ["f16", "f32", "f64", "f8E5M2", "f8E5M2FNUZ", "f8E4M3FN", "f8E4M3FNUZ"]
+_FLOAT_TYPES = [
+    "bf16",
+    "f16",
+    "f32",
+    "f64",
+    "f8E5M2",
+    "f8E5M2FNUZ",
+    "f8E4M3FN",
+    "f8E4M3FNUZ",
+]
 _INDEX_TYPES = ["index"]
 
 
@@ -55,9 +65,12 @@ def bitwidth(self):
             return 64
         if "f8" in self._name:
             return 8
+        if "bf16" in self._name:
+            return 16
         return int(self._name[1:])
 
 
+bf16 = DataType("bf16")
 bool = DataType("bool", "i1")
 i4 = DataType("i4")
 i8 = DataType("i8")

iree/turbine/kernel/lang/__init__.py

@@ -13,6 +13,7 @@
 )
 
 from .._support.dtype import (
+    bf16,
     bool,
     i4,
     i8,

iree/turbine/kernel/lang/global_symbols.py

@@ -1,16 +1,38 @@
 from .._support.indexing import index_symbol
+import sys
 
 # Global symbols used throughout the code.
 
 # Address spaces.
 GLOBAL_ADDRESS_SPACE = index_symbol("$GLOBAL_ADDRESS_SPACE")
 SHARED_ADDRESS_SPACE = index_symbol("$SHARED_ADDRESS_SPACE")
 
+
 # Distribution symbols.
 WORKGROUP_0 = index_symbol("$WG0")
 WORKGROUP_1 = index_symbol("$WG1")
 WORKGROUP_2 = index_symbol("$WG2")
 
+
+def create_additional_workgroup_symbols():
+    """
+    Since we can have a large number of workgroups, we create
+    symbols for them dynamically. However, only WORKGROUP_0,
+    WORKGROUP_1, and WORKGROUP_2 will persist during code generation,
+    so we generate those symbols statically.
+    """
+    max_workgroups = 5
+    for i in range(3, max_workgroups):
+        globals()[f"WORKGROUP_{i}"] = index_symbol(f"$WG{i}")
+
+
+def get_workgroup_symbol(i: int):
+    assert i >= 0, "Workgroup index must be non-negative."
+    return index_symbol(f"$WG{i}")
+
+
+create_additional_workgroup_symbols()
+
 THREAD_0 = index_symbol("$T0")
 THREAD_1 = index_symbol("$T1")
 THREAD_2 = index_symbol("$T2")

iree/turbine/kernel/wave/constraints.py

@@ -130,9 +130,17 @@ def mma_matrix_shapes(self, mma_type: Optional[MMAType]) -> tuple[int]:
                 return (16, 16, 16)
             case MMAType.F32_32x32x8_F16 | MMAType.I32_32x32x8_I8:
                 return (32, 32, 8)
-            case MMAType.F32_16x16x32_F8 | MMAType.F32_16x16x32_K4_F8 | MMAType.I32_16x16x32_I8:
+            case (
+                MMAType.F32_16x16x32_F8
+                | MMAType.F32_16x16x32_K4_F8
+                | MMAType.I32_16x16x32_I8
+            ):
                 return (16, 16, 32)
-            case MMAType.F32_32x32x16_F8 | MMAType.F32_32x32x16_K4_F8 | MMAType.I32_32x32x16_I8:
+            case (
+                MMAType.F32_32x32x16_F8
+                | MMAType.F32_32x32x16_K4_F8
+                | MMAType.I32_32x32x16_I8
+            ):
                 return (32, 32, 16)
             case _:
                 return ()
@@ -234,7 +242,11 @@ def apply(
                     1,  # N
                     1,  # K
                 ]
-            case MMAType.F32_16x16x32_F8 | MMAType.F32_16x16x32_K4_F8 | MMAType.I32_16x16x32_I8:
+            case (
+                MMAType.F32_16x16x32_F8
+                | MMAType.F32_16x16x32_K4_F8
+                | MMAType.I32_16x16x32_I8
+            ):
                 offset = [
                     Piecewise(
                         (lane % 16, ~MMA_ACC), (4 * floor(lane / 16), MMA_ACC)
@@ -262,7 +274,11 @@ def apply(
                         + 4 * floor(lane / 16)
                         + (GPR_NUM % 4),  # K
                     ]
-            case MMAType.F32_32x32x16_F8 | MMAType.F32_32x32x16_K4_F8 | MMAType.I32_32x32x16_I8:
+            case (
+                MMAType.F32_32x32x16_F8
+                | MMAType.F32_32x32x16_K4_F8
+                | MMAType.I32_32x32x16_I8
+            ):
                 offset = [
                     Piecewise(
                         (lane % 32, ~MMA_ACC),
@@ -328,6 +344,16 @@ class WorkgroupConstraint(Constraint):
     tile_size: IndexExpr
     workgroup_dim: int
 
+    def __post_init__(self):
+        self.wg_dim = None
+        match self.workgroup_dim:
+            case 0 | 1 | 2 | 3 | 4:
+                self.wg_dim = get_workgroup_symbol(self.workgroup_dim)
+            case _:
+                raise ValueError(
+                    "Invalid workgroup dimension. Expected 0, 1, 2, 3 or 4."
+                )
+
     @property
     def count(self) -> IndexExpr:
         """
@@ -336,16 +362,7 @@ def count(self) -> IndexExpr:
         return ceiling(self.dim / self.tile_size)
 
     def apply(self) -> IndexSequence:
-        match self.workgroup_dim:
-            case 0:
-                wg_dim = WORKGROUP_0
-            case 1:
-                wg_dim = WORKGROUP_1
-            case 2:
-                wg_dim = WORKGROUP_2
-            case _:
-                raise ValueError("Invalid workgroup dimension. Expected 0, 1 or 2.")
-        return IndexSequence(wg_dim * self.tile_size, 1)
+        return IndexSequence(self.wg_dim * self.tile_size, 1)
 
 
 def get_grid_shape(wg_constraints: list[WorkgroupConstraint]) -> list[IndexExpr]:

iree/turbine/kernel/wave/expansion.py

@@ -6,7 +6,7 @@
 
 import itertools
 import torch.fx as fx
-from typing import Any, TypeAlias
+from typing import Any, TypeAlias, Sequence, Type, Callable
 from functools import partial
 
 from .constraints import (
@@ -15,8 +15,23 @@
     WorkgroupConstraint,
     TilingConstraint,
 )
-from ..ops.wave_ops import *
-from .._support.indexing import IndexingContext, IndexSequence
+from ..ops.wave_ops import (
+    Allocate,
+    CustomOp,
+    GetResult,
+    Getitem,
+    IterArg,
+    MMA,
+    Output,
+    Placeholder,
+    Read,
+    ReduceOp,
+    Reduction,
+    Reshape,
+    Write,
+    get_custom,
+)
+from .._support.indexing import IndexingContext, IndexSymbol
 from ...support.logging import get_logger
 from .._support.tracing import CapturedTrace
 from .utils import (
@@ -440,7 +455,7 @@ def _expand_mma_reduction(
     for dim in mma.indexing_dims:
         if dim not in dim_scaling and mma.vector_shapes[dim] > 0:
             tile_size = idxc.get_static_value(dim)
-            dim_scaling[dim] = tile_size // mma.vector_shapes[dim]
+            dim_scaling[dim] = max(tile_size // mma.vector_shapes[dim], 1)
 
     # Store the original mma node and accumulator value for expansion.
     # When we begin expansion, we have a single mma node with the correct accumulator.