[TKW] Add missing validation and error messages (#432)

This is a collection of small changes adding additional validation or
more verbose error messages to TKW.

Includes:
- More verbose (or any) error messages for existing asserts and
exceptions. Maybe this is too much, but the current messages frequently
don't tell you what's going wrong. I think it's better to err on the
side of too much information. If these get annoying for some reason we
can trim them.
- Transforming some builtin `KeyError` into more specific error messages
with context
- Making the validation if `get_custom` is passed an op an error
instead. I think this more likely than not points to a bug, so it's
better for it to be an error. I can't remember the specific case in
which I hit it.
- Adding earlier validation if there's a type mismatch between a reduce
op init and return types.
- Adding earlier validation if there's an issue when decomposing reduce
ops and the local reduction doesn't match the accumulator reduction.
- Reporting the argument that has an issue if there's a failure in
decomposing reduce ops.
- Printing which node had an issue if there's a failure during codegen
- Validating that reduction and generated for loop have the same number
of arguments. This otherwise results in a failure later on, but we can
give more useful information here. I reported
#384 for the bug that
causes this to fire.
- Validating MMA shapes. `m` has to be in `lhs` and `n` in `rhs`.
Locally, I actually have much more restrictive validation that lhs had
to be `[..., m, k]` and rhs `[..., n, k]`. In theory it looks like Wave
is supposed to figure things out if that isn't the setup, but I never
had a case where it actually worked, so it seems like you need walk some
narrow path. This version is the less restrictive one though.
- Reporting more information if IREE invocation fails.
- Removing assumption that a reduction has users in `get_users`
- Removing some unused variables and arguments
- Adding a `__str__` method for IndexingContext and `__repr__` methods
to `ExpansionInfo` and `ReductionInfo`. Maybe these should be data
classes?
- Adding some missing types to some functions

One note is that I'm not really sure what the convention is for
Exception types in the project, so a lot of these are just RuntimeError.
That's not awesome, but I think it's still a lot more helpful than
nothing. I tried to avoid `raise ... from` as in my experience these
usually result in unhelpfully verbose stacks.

Author: GMNGeoffrey

iree/turbine/kernel/_support/indexing.py

@@ -117,6 +117,18 @@ def __init__(self):
         self.frozen_subs: list[tuple[IndexSymbol, int]] = []
         self.unbacked_symbols: list[IndexSymbol] = []
 
+    def __str__(self):
+        return (
+            f"IndexingContext("
+            f"subs: {self.subs}\n"
+            f"special_subs: {self.special_subs}\n"
+            f"shaped_bindings: {self.shaped_bindings}\n"
+            f"dyn_dims: {self.dyn_dims}\n"
+            f"frozen_subs: {self.frozen_subs}\n"
+            f"unbacked_symbols: {self.unbacked_symbols}\n"
+            ")"
+        )
+
     def next_dyn_dim(self) -> IndexSymbol:
         s = index_symbol(f"D{len(self.dyn_dims)}")
         self.dyn_dims.append(s)
@@ -157,7 +169,7 @@ def _bind_symbol(self, symbol: IndexSymbol, value: int):
         self.subs[symbol] = value
 
     def finalize(self):
-        assert len(self.frozen_subs) == 0
+        assert len(self.frozen_subs) == 0, f"{self.frozen_subs=}"
         # Go over everything we know and bind all free symbols.
         for _sb in self.shaped_bindings.values():
             for i in range(_sb.shaped_type.rank):

iree/turbine/kernel/compiler/kernel_codegen.py

@@ -346,7 +346,13 @@ def __init__(self, sig: KernelSignature, entry_block: Block):
         }
 
     def resolve_by_reference(self, reference: Any) -> Value:
-        binding = self._bindings_by_reference[reference]
+        try:
+            binding = self._bindings_by_reference[reference]
+        except KeyError:
+            pretty = "\n".join(
+                f"{k}: {v}" for k, v in self._bindings_by_reference.items()
+            )
+            raise KeyError(f"{reference} not in signature:\n{pretty}")
         return self.resolve(binding)
 
     @abstractmethod

iree/turbine/kernel/ops/wave_ops.py

@@ -354,8 +354,7 @@ def new_function(*args: Any, **kwargs: dict[str, Any]):
 def get_custom(node: fx.Node) -> "CustomOp":
     """Get the corresponding CustomOp for a given fx.Node."""
     if isinstance(node, CustomOp):
-        print("Careful! You passed a custom op where an fx.Node was required.")
-        return node
+        raise ValueError(f"fx.Node required but got custom op {node}")
     if not isinstance(node, fx.Node):
         raise ValueError(f"Expected an fx.Node but got {type(node)}")
 
@@ -549,7 +548,7 @@ def erase(self):
         self.graph.erase_node(self.fx_node)
 
     @classmethod
-    def handle(cls, graph, *args, **kwargs) -> fx.Node:
+    def handle(cls, graph: RegionGraph, *args, **kwargs) -> fx.Node:
         node = cls(*args, **kwargs)
         node._add_proxy_to_graph(graph)
         node.fx_node.node.tkw_op = cls
@@ -1407,7 +1406,12 @@ class Reduction(NestedRegionOp):
     implicit_captures: Sequence[fx.Proxy]
 
     @classmethod
-    def handle(cls, graph, *args, **kwargs):
+    def handle(cls, graph: RegionGraph, *args, **kwargs):
+        if not isinstance(graph, RegionGraph):
+            raise TypeError(
+                f"handle expected {RegionGraph.__name__} but got {type(graph)}"
+            )
+
         def wrapper(f):
             with graph.subtracer() as subtracer:
                 subgraph_name, implicit_captures = subtracer.trace(f)
@@ -1689,8 +1693,14 @@ class GetResult(CustomOp):
     res_idx: int
 
     def infer_type(self):
-        src_type = get_custom(self.value).type
+        op = get_custom(self.value)
+        src_type = op.type
         if isinstance(src_type, list):
+            if self.res_idx >= len(src_type):
+                raise RuntimeError(
+                    f"GetResult of {self.res_idx} from result with {len(src_type)} results"
+                    f"\n{op=}\nsrc={self.value}\n{src_type=}"
+                )
             self.type = src_type[self.res_idx]
         else:
             self.type = src_type
@@ -1703,7 +1713,7 @@ def indexing_dims(self) -> list[IndexExpr]:
         )
         src_indexing = get_custom(self.value).indexing_dims
         if has_multiple_value(src_indexing):
-            assert self.res_idx <= len(src_indexing) - 1
+            assert self.res_idx < len(src_indexing), f"{self=}"
             src_indexing = src_indexing[self.res_idx]
         assert is_valid_indexing_dim(src_indexing)
         return src_indexing
@@ -1715,11 +1725,11 @@ def index(self) -> dict[IndexSymbol, IndexSequence]:
         if custom_index is None:
             return None
         if not isinstance(custom, Reduction):
-            return custom.index
+            return custom_index
         assert isinstance(custom_index, Sequence) and self.res_idx < len(
             custom.indexing_dims
-        )
-        return custom.index[self.res_idx]
+        ), f"Invalid {custom_index=} with {self.res_idx=} and {custom.indexing_dims=}\n{custom}"
+        return custom_index[self.res_idx]
 
     @index.setter
     def index(self, value: dict[IndexSymbol, IndexSequence]):
@@ -1882,6 +1892,15 @@ def infer_type(self):
         reduced_dims = [dims for dims in src_type.symbolic_shape if dims != self.dim]
         dst_type = Register[(*reduced_dims, src_type.dtype)]
         self.type = dst_type
+        if (
+            self.init is not None
+            and get_custom(self.init).type.symbolic_shape != self.type.symbolic_shape
+        ):
+            raise RuntimeError(
+                f"Init type for {self.tkw_op_name} {get_custom(self.init).type.symbolic_shape}"
+                f" must match reduce type {self.type.symbolic_shape}"
+                f"\n{self}"
+            )
 
     @property
     def num_reduction_dims(self) -> int:

iree/turbine/kernel/wave/analysis/index_sequence_analysis.py

@@ -110,7 +110,7 @@ def verify_nodes(trace: CapturedTrace, constraints: list[Constraint]):
             continue
         if isinstance(custom, (Output, NestedRegionOp)):
             continue
-        assert custom.index, f"Index not set for node {custom.fx_node}"
+        assert custom.index, f"Index not set for node {custom.fx_node}: {custom}"
         if not custom.vector_shapes:
             # If vector_shapes is not set, see if it can be derived from the hardware constraints.
             hw_constraint = get_hardware_constraint(constraints)
@@ -121,7 +121,9 @@ def verify_nodes(trace: CapturedTrace, constraints: list[Constraint]):
                 custom.vector_shapes = {}
                 for dim in update_vector_shapes:
                     custom.vector_shapes[dim] = hw_constraint.vector_shapes[dim]
-        assert custom.vector_shapes, f"Vector shapes not set for node {custom.fx_node}"
+        assert (
+            custom.vector_shapes
+        ), f"Vector shapes not set for node {custom.fx_node}: {custom}"
 
 
 def set_node_indices(
@@ -685,7 +687,13 @@ def apply_offset(node: fx.Node):
             return False
         for dim, scale in custom.expanded_dims.items():
             if dim in custom.index:
-                custom.index[dim].start += scale * custom.vector_shapes[dim]
+                try:
+                    custom.index[dim].start += scale * custom.vector_shapes[dim]
+                except KeyError as e:
+                    raise RuntimeError(
+                        f"op index or vector shapes missing expanded dim {dim}:\n"
+                        f"{custom.index}\n{custom.vector_shapes}\n{custom}"
+                    )
         return False
 
     trace.walk(apply_offset)
@@ -741,8 +749,25 @@ def get_index(custom: CustomOp):
         lhs = get_custom(custom.lhs)
         rhs = get_custom(custom.rhs)
 
-        lhs_dim, lhs_size = get_largest_index_and_size(get_index(lhs))
-        rhs_dim, rhs_size = get_largest_index_and_size(get_index(rhs))
+        lhs_index = get_index(lhs)
+        rhs_index = get_index(rhs)
+
+        lhs_dim, lhs_size = get_largest_index_and_size(lhs_index)
+        rhs_dim, rhs_size = get_largest_index_and_size(rhs_index)
+
+        extra_error_info = (
+            f"\n{binary_op=}"
+            f"\n{lhs=}"
+            f"\n{lhs_index=}"
+            f"\n{lhs_dim=}"
+            f"\n{lhs_size=}"
+            f"\n{lhs.type.symbolic_shape=}"
+            f"\n{rhs=}"
+            f"\n{rhs_index=}"
+            f"\n{rhs_dim=}"
+            f"\n{rhs_size=}"
+            f"\n{rhs.type.symbolic_shape=}"
+        )
 
         # If they are equal we are done.
         if lhs_dim == rhs_dim and lhs_size == rhs_size:
@@ -753,7 +778,8 @@ def get_index(custom: CustomOp):
         # Cannot handle discrepancies when both shapes are > 1.
         if lhs_size > 1 and rhs_size > 1:
             raise NotImplementedError(
-                "Currently only support resolving discrepancies when one of the shapes is 1."
+                f"Currently only support resolving discrepancies when one of the shapes is 1."
+                f"{extra_error_info}"
             )
 
         broadcast_rhs = lhs_size > rhs_size
@@ -774,7 +800,9 @@ def get_index(custom: CustomOp):
 
             if not is_only_missing_dim and not is_innermost_dim:
                 raise NotImplementedError(
-                    "Currently only support resolving discrepancies when the broadcasting dimension is the innermost dimension."
+                    f"Currently only support resolving discrepancies when the broadcasting"
+                    f" dimension is the innermost dimension. {extra_error_info}"
+                    f"\n{broadcast_dim=}"
                 )
 
         # Broadcast

iree/turbine/kernel/wave/codegen/emitter.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, ClassVar, Optional, List, Type, Dict
 from dataclasses import dataclass
 from collections import namedtuple
+import sys
 
 import torch.fx as fx
 
@@ -112,7 +113,11 @@ def _emit_function_call_node(self, node: fx.Node):
         except KeyError:
             raise CodegenError(f"No handler registered for op {target_op}")
 
-        handler(self, node)
+        try:
+            handler(self, node)
+        except:
+            print(f"Error handling {node}", file=sys.stderr)
+            raise
 
     def lookup_node_values(self, node: fx.Node) -> List[Value]:
         assert NDEBUG or isinstance(node, fx.Node)

iree/turbine/kernel/wave/codegen/handlers.py

@@ -391,9 +391,15 @@ def handle_generic_binary(emitter: WaveEmitter, node: fx.Node):
             rhs = cast_py_value(emitter, rhs)
 
             if lhs.ir_value.type != rhs.ir_value.type:
+                op = get_custom(node)
                 raise ValidationError(
-                    "Expected lhs and rhs to have same type."
-                    f" Got: {lhs.ir_value.type} vs {rhs.ir_value.type}"
+                    f"Expected lhs and rhs to have same type for\n"
+                    f"{op}\nGot\n"
+                    f"lhs: {lhs.ir_value.type} vs rhs: {rhs.ir_value.type}\n"
+                    f"{lhs=}\n"
+                    f"{rhs=}\n"
+                    f"lhs={get_custom(op.lhs)}\n"
+                    f"rhs={get_custom(op.rhs)}"
                 )
 
             lhs = lhs.ir_value
@@ -768,6 +774,11 @@ def handle_reduction(emitter: WaveEmitter, node: fx.Node):
         # Add mapping for iter args.
         subgraph: fx.Graph = emitter.trace.get_subgraph(subgraph)
         iter_args: list[fx.Node] = get_custom(node).iter_args(subgraph)
+        assert len(iter_args) == len(forOp.inner_iter_args), (
+            f"Len of reduction and for op iter args must match,"
+            f" Reduction args: {iter_args};"
+            f" For Op args: {[a.type for a in forOp.inner_iter_args]}"
+        )
         for i, v in enumerate(forOp.inner_iter_args):
             emitter.bind_node_proxy(iter_args[i], IRProxyValue(v))
         captured_vars: list[fx.Node] = get_custom(node).captured_vars(subgraph)
@@ -785,6 +796,12 @@ def handle_reduction(emitter: WaveEmitter, node: fx.Node):
         flat_ret_values = [
             cast_py_value(emitter, value).ir_value for value in flat_ret_values
         ]
+        assert len(flat_ret_values) == len(flat_init_args), (
+            f"Loop must have the same number of return values as init args, but got\n"
+            f"{len(flat_ret_values)} vs {len(flat_init_args)}\n"
+            f"{flat_ret_values=}\n"
+            f"{flat_init_args=}\n"
+        )
         scf_d.YieldOp(flat_ret_values)
 
     emitter.bind_node_proxies(node, [IRProxyValue(v) for v in forOp.results_])
@@ -907,7 +924,7 @@ def handle_broadcast(emitter: WaveEmitter, node: fx.Node):
         raise NotImplementedError("Scalar src is not implemented yet for shuffleOp.")
     assert (
         vector_type.rank == 0 or vector_type.rank == 1
-    ), f"expected vector_type.rank == 1 but got {vector_type}"
+    ), f"expected vector_type.rank == 1 but got {vector_type}, {node}"
 
     # Handles scalar broadcast case.
     if vector_type.rank == 0:

iree/turbine/kernel/wave/decompose_reduce_ops.py

@@ -80,7 +80,7 @@ def determine_shuffle_config(
     return cluster_size, cluster_stride[0]
 
 
-def get_graph_node(custom: CustomOp, graph: fx.Graph):
+def get_graph_node(custom: CustomOp, graph: fx.Graph) -> fx.Node:
     custom.add_to_graph(graph)
     custom = custom.fx_node
     return custom
@@ -117,7 +117,7 @@ def emit_local_reduction(
     reduction_src: list[fx.Node],
     graph: fx.Graph,
     local_reduction_size,
-):
+) -> fx.Node:
     """
     Does reduction over all the element carried along by ReductionOp at local
     thread/SIMT level. This is done by reducing expanded sources combining them
@@ -135,7 +135,7 @@ def emit_scalarized_local_reduction(
     reduction_src: list[fx.Node],
     graph: fx.Graph,
     local_reduction_size,
-):
+) -> fx.Node:
     """
     Special case of local reduction wher we try to scalarize/get rid of most vector ops.
     this is useful for maximum, to expose more opportunities for v_max3_f32,
@@ -184,7 +184,6 @@ def emit_global_reduction(
 def decompose_reduce_ops(
     trace: CapturedTrace,
     constraints: list[Constraint],
-    index_map: dict[IndexSymbol, int],
 ):
     """
     The lowering for multi_reduction is done in two steps:
@@ -205,11 +204,6 @@ def decompose_reduce_ops(
     hardware_constraint = next(
         c for c in constraints if isinstance(c, HardwareConstraint)
     )
-    constraint_tile_size = {
-        c.dim: c.tile_size
-        for c in constraints
-        if isinstance(c, TilingConstraint) or isinstance(c, WorkgroupConstraint)
-    }
     induction_vars = [
         c.induction_var for c in constraints if isinstance(c, TilingConstraint)
     ]
@@ -242,9 +236,20 @@ def decompose_reduce_ops(
             get_thread_shape = lambda index: max(
                 subs_idxc(x.size) for x in index.values()
             )
-            local_reduce_sizes = [
-                get_thread_shape(get_custom(arg).index) for arg in reduction_src
-            ]
+            local_reduce_sizes = []
+            for arg in reduction_src:
+                try:
+                    op = get_custom(arg)
+
+                    thread_shape = get_thread_shape(op.index)
+                    local_reduce_sizes.append(thread_shape)
+                except Exception as e:
+                    index_str = "\n".join(f"{k}: {v}" for k, v in op.index.items())
+                    raise RuntimeError(
+                        f"Error in decompose_reduce_ops: {arg} with index\n"
+                        f"{index_str}\n{reduction_src=}\n{reduction_acc=}\n{reduction_dim=}"
+                    ) from e
+
             if not all_equal(local_reduce_sizes):
                 raise NotImplementedError(
                     "NYI: Expect all reduce_src to have same local reduce size."
@@ -258,6 +263,19 @@ def decompose_reduce_ops(
                     binary_fn, reduction_src, custom.graph, local_reduce_sizes[0]
                 )
 
+            if (
+                reduction_acc is not None
+                and get_custom(local_reduction).type.symbolic_shape
+                != get_custom(reduction_acc).type.symbolic_shape
+            ):
+                raise RuntimeError(
+                    "Local reduction and accumulator reduction must have same shape."
+                    f"\nlocal_reduction: {get_custom(local_reduction).type.symbolic_shape}"
+                    f"\nreduction_acc: {get_custom(reduction_acc).type.symbolic_shape}"
+                    f"\nlocal_reduction: {get_custom(local_reduction)}"
+                    f"\nreduction_acc: {get_custom(reduction_acc)}"
+                    f"\n{custom}"
+                )
             # Global Reduce
             cluster_size, cluster_stride = determine_shuffle_config(
                 reduction_src[0].index,