Assert whether tl.advance is unused at the AST level. (#6768)

Previously, we would check whether it is unused at the IR level, but
that has false positives: for example, DCE after resolving compile-time
branches and static_range expressions could result in the tl.advance
return value being discarded due to no error on the user's behalf.

We instead inspect the AST and see whether a tl.advance Call is the
child of an Expr node; if so, the user has almost certainly erred and we
should loudly communicate this.

Author: apgoucher

python/src/ir.cc

@@ -772,19 +772,7 @@ void init_triton_ir(py::module &&m) {
           },
           ret::reference)
       //  .def("has_attr", &::FuncOp::hasAttr)
-      .def("finalize",
-           [](FuncOp &self) -> void {
-             // Check if the result of tl.advance is used
-             self.walk([&](AdvanceOp op) {
-               if (op->getResult(0).use_empty())
-                 outputWarning(op->getLoc(), "The result of tl.advance is not "
-                                             "being used. Note that tl.advance "
-                                             "does not have any side effects. "
-                                             "To move the block pointer, you "
-                                             "need to assign the result of "
-                                             "tl.advance to a variable.");
-             });
-           })
+      .def("finalize", [](FuncOp &self) -> void {})
       .def_property_readonly("type", &FuncOp::getFunctionType)
       .def("reset_type", &FuncOp::setType);
 

python/test/unit/language/test_compile_errors.py

@@ -453,3 +453,35 @@ def dot_kernel():
         assert (str(e.value.__cause__) == "max_num_imprecise_acc (128) must be <= K (64)")
     except AssertionError as assertion_err:
         raise assertion_err from e.value
+
+
+extra_words = "These are extra words in the error message."
+
+
+@triton.must_use_result(extra_words)
+@triton.jit
+def cube(x):
+    return x * x * x
+
+
+def test_unused_result():
+
+    @triton.jit
+    def evil_cube_kernel():
+        a = tl.full((64, 64), 0.0, tl.float32)
+        cube(a)
+
+    @triton.jit
+    def good_cube_kernel():
+        a = tl.full((64, 64), 0.0, tl.float32)
+        a = cube(a)
+
+    triton.compile(triton.compiler.ASTSource(fn=good_cube_kernel, signature={}, constexprs={}))
+
+    with pytest.raises(CompilationError) as e:
+        triton.compile(triton.compiler.ASTSource(fn=evil_cube_kernel, signature={}, constexprs={}))
+
+    expected_err_msg = "The result of cube is not being used. " + extra_words
+    obtained_err_msg = str(e.value).split('\n')[-1]
+
+    assert expected_err_msg == obtained_err_msg

python/triton/__init__.py

@@ -26,6 +26,8 @@
 from . import testing
 from . import tools
 
+must_use_result = language.core.must_use_result
+
 __all__ = [
     "autotune",
     "cdiv",
@@ -39,6 +41,7 @@
     "KernelInterface",
     "language",
     "MockTensor",
+    "must_use_result",
     "next_power_of_2",
     "OutOfResources",
     "reinterpret",

python/triton/compiler/code_generator.py

@@ -1220,6 +1220,13 @@ def visit_Call(self, node):
         if static_implementation is not None:
             return static_implementation(self, node)
 
+        mur = getattr(fn, '_must_use_result', False)
+        if mur and getattr(node, '_is_unused', False):
+            error_message = ["The result of %s is not being used." % ast.unparse(node.func)]
+            if isinstance(mur, str):
+                error_message.append(mur)
+            raise CompilationError(self.jit_fn.src, node, " ".join(error_message))
+
         kws = dict(self.visit(keyword) for keyword in node.keywords)
         args = [self.visit(arg) for arg in node.args]
         args = list(itertools.chain.from_iterable(x if isinstance(x, list) else [x] for x in args))
@@ -1277,6 +1284,7 @@ def visit_Attribute(self, node):
         return getattr(lhs, node.attr)
 
     def visit_Expr(self, node):
+        node.value._is_unused = True
         ast.NodeVisitor.generic_visit(self, node)
 
     def visit_NoneType(self, node):

python/triton/language/core.py

@@ -23,6 +23,14 @@
 PropagateNan = ir.PROPAGATE_NAN
 
 
+def must_use_result(x, s=True):
+    """If the result of this function is unused, throw an error."""
+    if isinstance(x, str):
+        return (lambda fn: must_use_result(fn, x))
+    x._must_use_result = s
+    return x
+
+
 def builtin(fn: T) -> T:
     """Mark a function as a builtin."""
     assert callable(fn)
@@ -2080,6 +2088,9 @@ def make_block_ptr(base: tensor, shape, strides, offsets, block_shape, order, _b
     return semantic.make_block_ptr(base, shape, strides, offsets, block_shape, order, _builder)
 
 
+@must_use_result(
+    "Note that tl.advance does not have any side effects. To move the block pointer, you need to assign the result of tl.advance to a variable."
+)
 @_tensor_member_fn
 @builtin
 def advance(base, offsets, _builder=None):