[TK] Add a TestLaunchContext for generating test dispatch IRs (#429)

This patch changes the way we do grid bindings. It is now expected that
the IndexingContext will have the workload values and the grid on
construction will try to build it's dims using them. Currently, all
symbols are assumed to have a constant value, so the grid is able to map
workgroup calculations to constants.

Ideally, we would have support to parameterize the kernel by some
symbols, which would be passed as workload bindings (We have no support
for this currently).

Author: Groverkss

core/shark_turbine/kernel/_support/indexing.py

@@ -104,13 +104,6 @@ def __new__(
         new_class.__qualname__ = repr(new_class)
         return new_class
 
-    def __class_getitem__(
-        cls, symbolic_shape: Union[SymbolicDimable, tuple[SymbolicShapeable]]
-    ) -> Type["Grid"]:
-        if not isinstance(symbolic_shape, tuple):
-            symbolic_shape = (symbolic_shape,)
-        return cast(Grid, _make_shaped_grid(cls, make_symbolic_shape(symbolic_shape)))
-
     def __repr__(self):
         if self.symbolic_shape:
             return f"Grid[{', '.join(repr(s) for s in self.symbolic_shape)}]"
@@ -122,20 +115,31 @@ class Grid(metaclass=_GridMeta, symbolic_shape=None):
     """Grid with bounding symbolic shape information in the type."""
 
     symbolic_shape: ClassVar[Optional[SymbolicShapeExpr]]
+    # TODO: dims should also allow dynamic dimensions.
+    dims: list[int]
     rank: int
 
-    def __init__(self, *dims: int):
-        rank = len(dims)
-        if self.symbolic_shape is not None:
-            if rank != len(self.symbolic_shape):
-                raise ValueError(
-                    f"Cannot create {type(self)}({', '.join(str(i) for i in dims)}): mismatched symbolic rank"
-                )
+    def __init__(self):
+        # Resolve the symbolic shape to concrete values.
+        idxc = IndexingContext.current()
+        if self.symbolic_shape:
+            dims = [idxc.get_static_value(dim) for dim in self.symbolic_shape]
+            if None in dims:
+                raise ValueError(f"NYI: Dynamic dims in Grid")
+            self.dims = cast(list[int], dims)
+        else:
+            self.dims = []
 
-        self.dims = dims
         # Shadow the type rank with the actual, which makes it concrete
         # for the generic case.
-        self.rank = rank
+        self.rank = len(self.dims)
+
+    def __class_getitem__(
+        cls, symbolic_shape: Union[SymbolicDimable, tuple[SymbolicShapeable]]
+    ) -> Type["Grid"]:
+        if not isinstance(symbolic_shape, tuple):
+            symbolic_shape = (symbolic_shape,)
+        return cast(Grid, _make_shaped_grid(cls, make_symbolic_shape(symbolic_shape)))
 
     def __repr__(self):
         return f"{repr(type(self))}({', '.join(str(i) for i in self.dims)})"
@@ -161,6 +165,8 @@ class ShapedGrid(Grid, symbolic_shape=symbolic_shape):
 # KernelBuffer
 ###############################################################################
 
+Dims = list[Union[None, IndexSymbol, int]]
+
 
 class KernelBufferUsage(Enum):
     NONE = 0
@@ -331,7 +337,6 @@ class TemporaryBuffer(KernelBuffer):
 ###############################################################################
 
 ShapedType = Union[Type[KernelBuffer], Type[Grid]]
-Dims = list[Union[None, IndexSymbol, int]]
 
 
 @dataclass(slots=True)

core/shark_turbine/kernel/_support/tracing.py

@@ -23,9 +23,11 @@
     backed_sym_index_type,
     BoundedRelation,
     IndexExpr,
+    IndexSymbol,
     Grid,
     KernelBuffer,
     SymIndex,
+    IndexingContext,
 )
 
 from ..lang.types import (
@@ -388,10 +390,16 @@ def __call__(self, *args, **kwargs):
     def eager_execute(self, args, kwargs):
         ...
 
+    def test_execute(self, args, kwargs):
+        ...
+
 
 class LaunchContext(ABC):
     __tk_context_idname__ = "ExecutionContext"
 
+    def __init__(self, constant_bindings: Dict[IndexSymbol, int] = {}):
+        self.constant_bindings = constant_bindings
+
     @staticmethod
     def current() -> "LaunchContext":
         try:
@@ -404,9 +412,21 @@ def current() -> "LaunchContext":
             return DebugLaunchContext()
 
     def __enter__(self) -> "LaunchContext":
+        # Push an indexing context with the constand bindings for this launch
+        # context in it.
+        # TODO: Is creating a IndexingContext as part of LaunchContext the
+        # correct layering?
+        idxc = IndexingContext()
+        context.push(IndexingContext, idxc)
+        for s, val in self.constant_bindings.items():
+            idxc.bind_constant(s, val)
         return context.push(LaunchContext, self)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        # Pop the indexing context created as part of this launch.
+        # TODO: Is creating a IndexingContext as part of LaunchContext the
+        # correct layering?
+        context.pop(IndexingContext, IndexingContext().current())
         context.pop(LaunchContext, self)
 
     @abstractmethod
@@ -419,6 +439,11 @@ def launch(self, launchable: Launchable, args, kwargs):
         return launchable.eager_execute(args, kwargs)
 
 
+class TestLaunchContext(LaunchContext):
+    def launch(self, launchable: Launchable, args, kwargs):
+        return launchable.test_execute(args, kwargs)
+
+
 ###############################################################################
 # Helpers
 ###############################################################################

core/shark_turbine/kernel/compiler/dispatch_codegen.py

@@ -5,7 +5,7 @@
 embedding and generating the calls/dispatches.
 """
 
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Type
 
 from .._support.indexing import (
     IndexingContext,
@@ -44,6 +44,8 @@
     KernelSignature,
 )
 
+from .._support.indexing import Grid
+
 
 class StreamExecutable:
     """Encapsulates a 'stream' compilable executable which can be dispatched to.
@@ -93,6 +95,7 @@ def define_entrypoint(
         self,
         name: str,
         sig: KernelSignature,
+        grid: Grid,
     ) -> "DispatchEntrypoint":
         """Defines a dispatch function with a signature like:
 
@@ -105,7 +108,7 @@ def define_entrypoint(
         Also adds an export with workgroup function like:
 
         ```
-        stream.executable.export public @name(%workload0 : index, %workload1 : index) -> (index, [[grid_arity...]]) {
+        stream.executable.export private @name(%workload0 : index, %workload1 : index) -> (index, [[grid_arity...]]) {
 
         }
         ```
@@ -115,28 +118,32 @@ def define_entrypoint(
         kb_input_bindings = sig.kernel_buffer_input_bindings
         kb_temp_bindings = sig.kernel_buffer_temporary_bindings
         kb_output_bindings = sig.kernel_buffer_output_bindings
-        # TODO: The way we are doing grid bindings is wrong. The Grid type should be paramerized
-        # with special grid axis symbols which are algebraically related to concrete shape dim
-        # symbols. For now, we are just treating the grid symbol as the input and output to the
-        # workload function, when in reality, the workload needs to derive from its leaf inputs.
-        grid_axis_bindings = sig.grid_bindings
+        # TODO: The way we are doing grid bindings is wrong. The Grid type
+        # should be paramerized with special grid axis symbols which are
+        # algebraically related to concrete shape dim symbols. For now, we are
+        # just assuming that the grid dims can be resolved to constants , when
+        # in reality, we should pass the workload and parameterize the grid
+        # dims on the workloads.
+        workload_axis_bindings = []
 
         # Input bindings are always user specified.
         # Grid/workgroup bindings are in the inputs section but are implied.
         # Temp bindings are a special kind of output bindings.
         # Output bindings are the real outputs.
         linear_bindings = (
             kb_input_bindings
-            + grid_axis_bindings
+            + workload_axis_bindings
             + kb_temp_bindings
             + kb_output_bindings
         )
 
-        # TODO: This is sloppy. This assert will hit on some user errors for unsupported
-        # type combinations and is just a last resort right now.
-        assert len(linear_bindings) == len(
-            sig.bindings
-        ), f"Not all bindings converted: {linear_bindings} vs {sig.bindings}"
+        # TODO: This is sloppy. This assert will hit on some user errors for
+        # unsupported type combinations and is just a last resort right now.
+        # TODO: This is currently disabled because the grid_bindings don't match
+        # workload bindings.
+        # assert len(linear_bindings) == len(
+        #     sig.bindings
+        # ), f"Not all bindings converted: {linear_bindings} vs {sig.bindings}"
 
         with self._loc:
             binding_type = IrType.parse("!stream.binding")
@@ -161,17 +168,22 @@ def abi_type(binding: BindingDesc):
             with InsertionPoint.at_block_begin(self._exe_block):
                 export_op = stream_d.ExecutableExportOp(name, name)
                 export_block = export_op.workgroup_count.blocks.append(
-                    *([b.as_mlir_type() for b in grid_axis_bindings])
+                    *([b.as_mlir_type() for b in workload_axis_bindings])
                 )
 
-            # TODO: Reify actual workload calculation.
             workgroup_builder = WorkgroupBuilder(
                 export_block, lambda vs: stream_d.ReturnOp(vs)
             )
-            workgroup_values = list(workgroup_builder.workload)
-            while len(workgroup_values) < 3:
-                with InsertionPoint(workgroup_builder.entry_block):
-                    result_type = IndexType.get()
+
+            # TODO: Support passing workload to the dispatch function.
+            with InsertionPoint(workgroup_builder.entry_block):
+                result_type = IndexType.get()
+                workgroup_values = [
+                    arith_d.constant(result_type, IntegerAttr.get(result_type, dim))
+                    for dim in grid.dims
+                ]
+
+                while len(workgroup_values) < 3:
                     workgroup_values.append(
                         arith_d.constant(result_type, IntegerAttr.get(result_type, 1))
                     )
@@ -220,8 +232,7 @@ def __init__(
     def resolve(self, binding: BindingDesc) -> Value:
         ref_type, ref_value = binding.reference
         if ref_type == "grid":
-            # TODO: Switch to stream op when #15889 is landed.
-            return flow_d.dispatch_workgroup_id(
+            return stream_d.dispatch_workgroup_id(
                 IntegerAttr.get(IndexType.get(), ref_value)
             )
 

core/shark_turbine/kernel/compiler/host_codegen.py

@@ -0,0 +1,58 @@
+from .kernel_codegen import KernelSignature
+from .dispatch_codegen import StreamExecutable
+
+from .builder import (
+    ModuleBuilder,
+)
+
+from .ir import (
+    Block,
+    FunctionType,
+    InsertionPoint,
+    IrType,
+    Location,
+    ArrayAttr,
+    SymbolRefAttr,
+    MemRefType,
+    RankedTensorType,
+    flow_d,
+    func_d,
+)
+
+
+def memref_to_tensor(memrefs: list[IrType]):
+    tensors = []
+    for m in memrefs:
+        assert isinstance(m, MemRefType)
+        t = RankedTensorType.get(m.shape, m.element_type)
+        tensors.append(t)
+    return tensors
+
+
+def isolated_test_call(
+    mb: ModuleBuilder, exe: StreamExecutable, sig: KernelSignature, entrypoint: str
+):
+    with InsertionPoint(mb.body_block), Location.unknown():
+        input_types = [b.as_mlir_type() for b in sig.kernel_buffer_input_bindings]
+        input_tensors = memref_to_tensor(input_types)
+        output_types = [b.as_mlir_type() for b in sig.kernel_buffer_output_bindings]
+        output_tensors = memref_to_tensor(output_types)
+
+        ftype = FunctionType.get(input_tensors, output_tensors)
+        func_op = func_d.FuncOp("isolated_benchmark", ftype)
+        arg_locs = [
+            (Location.name(b.name) if b.name is not None else Location.unknown())
+            for b in sig.kernel_buffer_input_bindings
+        ]
+        entry_block = func_op.add_entry_block(arg_locs)
+        with InsertionPoint(entry_block):
+            assert isinstance(entry_block, Block)
+            # Create a flow.dispatch op to the kernel
+            dispatch = SymbolRefAttr.get([exe.sym_name.value, entrypoint])
+            entrypoints = ArrayAttr.get([dispatch])
+
+            out = flow_d.DispatchOp(
+                output_tensors, [], entrypoints, entry_block.arguments, [], []
+            )
+
+            func_d.ReturnOp(out)

core/shark_turbine/kernel/compiler/ir.py

@@ -2,8 +2,11 @@
     AffineConstantExpr,
     AffineExpr,
     AffineMap,
+    FlatSymbolRefAttr,
+    SymbolRefAttr,
     AffineMapAttr,
     Attribute,
+    RankedTensorType,
     ArrayAttr,
     Block,
     Context,

core/shark_turbine/kernel/compiler/kernel_codegen.py

@@ -118,7 +118,6 @@ def grid_bindings(self) -> list[BindingDesc]:
     @property
     def kernel_buffer_input_bindings(self) -> list[BindingDesc]:
         """Gets all kernel buffer bindings with input usage."""
-        print("ALL=", self.bindings)
         return [
             b
             for b in self.bindings

core/shark_turbine/kernel/gen/__init__.py

@@ -1 +1,3 @@
 from .thread import *
+
+from .._support.tracing import TestLaunchContext