Add mask to tl.device_assert (triton-lang#7905)

<!---
The core Triton is a small number of people, and we receive many PRs
(thank
you!).  To help us review your code more quickly, **if you are a new
contributor (less than 3 PRs merged) we ask that you complete the
following
tasks and include the filled-out checklist in your PR description.**

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->


- [x] I am not making a trivial change, such as fixing a typo in a
comment.

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
  - [x] I have added tests.

Add `mask=mask` to `tl.device_assert`

Frequently in Triton kernels, data is loaded with a mask. The data
outside the mask has unknown value. This makes it frustrating to write
`tl.device_assert`s. In particular, I have to use this pattern:
`tl.device_assert((data == 1337) | ~mask)`. This works fine; it's just a
little unpleasant, and you have to remember that `==` has lower
precedence than `|`. With this PR, you can instead write
`tl.device_assert(data == 1337, mask=mask)`.

Representative example:

```py
import torch
import triton.language as tl
import triton
torch.set_default_device("cuda")

@triton.jit
def kernel(data_ptr, len):
	BLOCK_SZ: tl.constexpr = 32

	mask = tl.arange(0, BLOCK_SZ) < len
	data = tl.load(data_ptr + tl.arange(0, BLOCK_SZ), mask=mask)
	tl.device_assert(data == 1337, mask=mask)

tensor = torch.full((10,), 1337, dtype=torch.int32)
print("succeeds:")
kernel[(1,)](tensor, len(tensor))
print("fails:")
kernel[(1,)](tensor, len(tensor) + 1) # the next int is probably not 1337
```

Author: leijurv

python/test/unit/test_debug.py

@@ -5,31 +5,32 @@
 
 
 @pytest.mark.parametrize('cond', [True, False])
+@pytest.mark.parametrize('mask', [True, False, None])
 @pytest.mark.parametrize('opt_flag', [True, False, None])
 @pytest.mark.parametrize('env_var', [True, False])
 @pytest.mark.parametrize('jit_flag', [True, False])
 @pytest.mark.forked
-def test_device_assert(monkeypatch, cond, opt_flag, env_var, jit_flag, device):
+def test_device_assert(monkeypatch, cond, mask, opt_flag, env_var, jit_flag, device):
     monkeypatch.setenv("TRITON_DEBUG", str(int(env_var)))
     torch.zeros([1], dtype=torch.int32, device=device)
 
     @triton.jit(debug=jit_flag)
-    def _kernel(COND: tl.constexpr):
-        tl.device_assert(COND, 'test')
+    def _kernel(COND: tl.constexpr, MASK: tl.constexpr):
+        tl.device_assert(COND, 'test', mask=MASK)
 
     is_debug = env_var or (opt_flag if opt_flag is not None else jit_flag)
 
     kwargs = {}
     if opt_flag is not None:
         kwargs["debug"] = opt_flag
 
-    if not cond and is_debug:
+    if not cond and is_debug and mask is not False:
         with pytest.raises(RuntimeError):
-            _kernel[(1, )](cond, **kwargs)
+            _kernel[(1, )](cond, mask, **kwargs)
             getattr(torch, device).synchronize()
         return
 
-    _kernel[(1, )](cond, **kwargs)
+    _kernel[(1, )](cond, mask, **kwargs)
     getattr(torch, device).synchronize()
 
 

python/triton/language/core.py

@@ -2984,7 +2984,7 @@ def device_print(prefix, *args, hex=False, _semantic=None):
 
 
 @builtin
-def device_assert(cond, msg="", _semantic=None):
+def device_assert(cond, msg="", mask=None, _semantic=None):
     '''
     Assert the condition at runtime from the device.  Requires that the environment variable :code:`TRITON_DEBUG`
     is set to a value besides :code:`0` in order for this to have any effect.
@@ -3003,7 +3003,10 @@ def device_assert(cond, msg="", _semantic=None):
     :param msg: the message to print if the assertion fails. This is required to be a string literal.
     '''
     msg = _unwrap_if_constexpr(msg)
-    return _semantic.device_assert(_semantic.to_tensor(cond), msg)
+    mask = _unwrap_if_constexpr(mask)
+    if mask is not None:
+        mask = _semantic.to_tensor(mask)
+    return _semantic.device_assert(_semantic.to_tensor(cond), msg, mask)
 
 
 @builtin

python/triton/language/semantic.py

@@ -219,7 +219,7 @@ def binary_op_sanitize_overflow_impl(self, lhs: TensorTy, rhs: TensorTy, binary_
         min_value = self.scalar_constant(min_value, tl.int64)
         cond = self.and_(self.less_equal(ret, max_value), self.greater_equal(ret, min_value))
         msg = f"int{lhs_sca_ty.int_bitwidth} overflow detected for operation {binary_op.__name__}"
-        self.device_assert(cond, msg)
+        self.device_assert(cond, msg, None)
 
     def add(self, input: TensorTy | numbers.Number, other: TensorTy | numbers.Number,
             sanitize_overflow: bool) -> TensorTy:
@@ -1808,9 +1808,11 @@ def device_print(self, prefix: str, args: List[TensorTy], hex: bool) -> TensorTy
         is_signed = [arg.dtype.is_int_signed() for arg in args]
         return self.tensor(self.builder.create_print(prefix, hex, new_args, is_signed), tl.void)
 
-    def device_assert(self, cond: TensorTy, msg: str) -> TensorTy:
+    def device_assert(self, cond: TensorTy, msg: str, mask: Optional[TensorTy]) -> TensorTy:
         if not self.builder.options.debug:
             return
+        if mask is not None:
+            cond = self.or_(cond, self.not_(mask))
         return self.tensor(self.builder.create_assert(cond.handle, msg), tl.void)
 
     def assume(self, cond) -> TensorTy: