Skip to content

Commit d92727c

Browse files
authored
[FRONTEND] Add descriptive messages to bare asserts (triton-lang#10405)
## Summary Improve error diagnostics for user-facing `assert` statements in `semantic.py` and `core.py` by adding descriptive messages following the "expected X, got Y" pattern. As requested by @lezcano in triton-lang#10341 — previously, users hitting these asserts (e.g., mismatched accumulator dtype in `tl.dot`) would get an empty `AssertionError` with no guidance on what went wrong. ## Changes - `semantic.py`: 17 bare asserts → descriptive messages - `core.py`: 5 bare asserts → descriptive messages - Split compound `assert shape == X and dtype == Y` into two separate asserts for clearer diagnostics
1 parent 6650ee3 commit d92727c

2 files changed

Lines changed: 37 additions & 22 deletions

File tree

python/triton/language/core.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3250,7 +3250,7 @@ def selu(x, alpha):
32503250
:return: one tensor or a tuple of tensors, depending on the mapped function.
32513251
'''
32523252
# Build the block for the nested region first to discover the return types
3253-
assert pack >= 1
3253+
assert pack >= 1, f"pack must be >= 1, got {pack}"
32543254
in_scalar_tys = [t.type.scalar for t in args]
32553255
builder = _semantic.builder
32563256
block = builder.new_block()
@@ -3882,8 +3882,8 @@ def builtin_max(*args, propagate_nan=_NOTHING, _semantic=None):
38823882
is_constexpr = all(not isinstance(x, base_value) for x in args)
38833883
if is_constexpr:
38843884
assert propagate_nan is _NOTHING, "propagate_nan is not supported on builtin max"
3885-
assert not any(math.isnan(x) for x in args)
3886-
assert not any(is_negative_zero(x) for x in args)
3885+
assert not any(math.isnan(x) for x in args), "constexpr max does not support NaN values"
3886+
assert not any(is_negative_zero(x) for x in args), "constexpr max does not support negative zero"
38873887
return constexpr(builtins.max(_unwrap_if_constexpr(args)))
38883888

38893889
if propagate_nan is _NOTHING:
@@ -3906,8 +3906,8 @@ def builtin_min(*args, propagate_nan=_NOTHING, _semantic=None):
39063906
is_constexpr = all(not isinstance(x, base_value) for x in args)
39073907
if is_constexpr:
39083908
assert propagate_nan is _NOTHING, "propagate_nan is not supported on builtin min"
3909-
assert not any(math.isnan(x) for x in args)
3910-
assert not any(is_negative_zero(x) for x in args)
3909+
assert not any(math.isnan(x) for x in args), "constexpr min does not support NaN values"
3910+
assert not any(is_negative_zero(x) for x in args), "constexpr min does not support negative zero"
39113911
return constexpr(builtins.min(_unwrap_if_constexpr(args)))
39123912

39133913
if propagate_nan is _NOTHING:

python/triton/language/semantic.py

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -214,8 +214,8 @@ def binary_op_sanitize_overflow_impl(self, lhs: TensorTy, rhs: TensorTy, binary_
214214
return
215215
lhs_sca_ty = lhs.type.scalar
216216
rhs_sca_ty = rhs.type.scalar
217-
assert lhs_sca_ty == rhs_sca_ty
218-
assert lhs_sca_ty.is_int()
217+
assert lhs_sca_ty == rhs_sca_ty, f"expected matching operand types, got {lhs_sca_ty} and {rhs_sca_ty}"
218+
assert lhs_sca_ty.is_int(), f"expected integer type, got {lhs_sca_ty}"
219219
lhs = self.cast(lhs, tl.int64)
220220
rhs = self.cast(rhs, tl.int64)
221221
ret = binary_op(lhs, rhs, False)
@@ -657,7 +657,7 @@ def expand_dims(self, input: TensorTy, axis: int) -> TensorTy:
657657

658658
def cat(self, lhs: TensorTy, rhs: TensorTy, can_reorder: bool) -> TensorTy:
659659
assert can_reorder, "current implementation of `cat` always may reorder elements"
660-
assert len(lhs.shape) == 1
660+
assert len(lhs.shape) == 1, f"expected 1D input for cat, got {len(lhs.shape)}D"
661661
ret_type = tl.block_type(lhs.type.scalar, [lhs.shape[0] + rhs.shape[0]])
662662
return self.tensor(self.builder.create_cat(lhs.handle, rhs.handle), ret_type)
663663

@@ -686,8 +686,9 @@ def join(self, a: TensorTy, b: TensorTy) -> TensorTy:
686686
return ret
687687

688688
def split(self, a: TensorTy) -> Tuple[TensorTy, TensorTy]:
689-
assert (len(a.shape) > 0)
690-
assert (tl._unwrap_if_constexpr(a.shape[-1]) == 2)
689+
assert (len(a.shape) > 0), "split requires a non-scalar tensor"
690+
assert (tl._unwrap_if_constexpr(a.shape[-1]) == 2), \
691+
f"expected last dimension to be 2 for split, got {tl._unwrap_if_constexpr(a.shape[-1])}"
691692

692693
new_shape = a.shape[:-1]
693694
ret_type = tl.block_type(a.type.scalar, new_shape)
@@ -754,7 +755,8 @@ def broadcast_impl_value(self, lhs: TensorTy, rhs: TensorTy) -> TensorTy:
754755
tl.block_type(rhs_ty.scalar, [1] + rhs_shape.values))
755756
rhs_ty = rhs.type
756757
rhs_shape = rhs_ty.get_block_shapes()
757-
assert len(rhs_shape) == len(lhs_shape)
758+
assert len(rhs_shape) == len(lhs_shape), \
759+
f"expected tensors of equal rank for broadcast, got {len(lhs_shape)} and {len(rhs_shape)}"
758760

759761
ret_shape = []
760762
for i, left in enumerate(lhs_shape):
@@ -1062,7 +1064,8 @@ def load(self, ptr: TensorTy, mask: Optional[TensorTy], other: Optional[TensorTy
10621064

10631065
def descriptor_load(self, desc: tl.tensor_descriptor_base, offsets, cache_modifier: str,
10641066
eviction_policy: str) -> TensorTy:
1065-
assert isinstance(desc, tl.tensor_descriptor_base)
1067+
assert isinstance(desc, tl.tensor_descriptor_base), \
1068+
f"expected a tensor descriptor, got {type(desc).__name__}"
10661069
ndim = len(desc.block_shape)
10671070
assert len(offsets) == ndim, f"expected {ndim} offsets, but got {len(offsets)}"
10681071

@@ -1072,10 +1075,12 @@ def descriptor_load(self, desc: tl.tensor_descriptor_base, offsets, cache_modifi
10721075
return self.tensor(x, desc.block_type)
10731076

10741077
def validate_store_like(self, desc: tl.tensor_descriptor_base, value: TensorTy, offsets) -> None:
1075-
assert isinstance(desc, tl.tensor_descriptor_base)
1078+
assert isinstance(desc, tl.tensor_descriptor_base), \
1079+
f"expected a tensor descriptor, got {type(desc).__name__}"
10761080
ndim = len(desc.block_shape)
10771081
assert len(offsets) == ndim, f"expected {ndim} offsets, but got {len(offsets)}"
1078-
assert value.shape == desc.block_shape
1082+
assert value.shape == desc.block_shape, \
1083+
f"expected value shape {desc.block_shape}, got {value.shape}"
10791084

10801085
def descriptor_store(self, desc: tl.tensor_descriptor_base, value: TensorTy, offsets) -> TensorTy:
10811086
self.validate_store_like(desc, value, offsets)
@@ -1136,7 +1141,8 @@ def descriptor_atomic_xor(self, desc: tl.tensor_descriptor_base, value: TensorTy
11361141
return self.tensor(self.builder.create_descriptor_reduce(kind, desc.handle, value.handle, offsets), tl.void)
11371142

11381143
def descriptor_gather(self, desc, x_offsets, y_offset, cache_modifier: str, eviction_policy: str) -> TensorTy:
1139-
assert isinstance(desc, tl.tensor_descriptor_base)
1144+
assert isinstance(desc, tl.tensor_descriptor_base), \
1145+
f"expected a tensor descriptor, got {desc.__class__.__name__}"
11401146
assert cache_modifier == "", "cache modifier is not supported yet"
11411147
assert eviction_policy == "", "eviction policy is not supported yet"
11421148

@@ -1162,7 +1168,8 @@ def descriptor_gather(self, desc, x_offsets, y_offset, cache_modifier: str, evic
11621168
return self.tensor(x, type)
11631169

11641170
def descriptor_scatter(self, desc, value: TensorTy, x_offsets, y_offset) -> TensorTy:
1165-
assert isinstance(desc, tl.tensor_descriptor_base)
1171+
assert isinstance(desc, tl.tensor_descriptor_base), \
1172+
f"expected a tensor descriptor, got {type(desc).__name__}"
11661173

11671174
# Validate descriptor.
11681175
assert len(desc.block_shape) == 2, f"descriptor must be 2D, but got {desc.block_shape}"
@@ -1421,7 +1428,7 @@ def _str_to_dot_input_precision(self, input_precision):
14211428

14221429
def dot(self, lhs: TensorTy, rhs: TensorTy, acc: TensorTy, input_precision: Optional[str],
14231430
max_num_imprecise_acc: int, out_dtype: tl.dtype | None) -> TensorTy:
1424-
assert lhs.type.is_block() and rhs.type.is_block()
1431+
assert lhs.type.is_block() and rhs.type.is_block(), "dot operands must be block tensors (not scalars)"
14251432

14261433
if lhs.dtype.is_fp8() and rhs.dtype.is_fp8():
14271434
# All combinations of supported fp8 x fp8 are permitted
@@ -1500,7 +1507,11 @@ def dot(self, lhs: TensorTy, rhs: TensorTy, acc: TensorTy, input_precision: Opti
15001507
acc_handle = self.builder.create_splat(ret_ty.to_ir(self.builder), _0)
15011508
else:
15021509
acc_handle = acc.handle
1503-
assert acc.type.shape == ret_ty.shape and acc.type.element_ty == out_dtype
1510+
assert acc.type.shape == ret_ty.shape, \
1511+
f"expected accumulator shape {ret_ty.shape}, got {acc.type.shape}"
1512+
assert acc.type.element_ty == out_dtype, \
1513+
f"expected accumulator dtype {out_dtype}, got {acc.type.element_ty}; " \
1514+
f"pass out_dtype={acc.type.element_ty} to use this accumulator dtype"
15041515

15051516
# max_num_imprecise_acc only applies to fp8 -> fp32 dot on sm_90
15061517
if max_num_imprecise_acc is None:
@@ -1565,7 +1576,7 @@ def verify_scaled_shape(self, M, N, K, lhs_scale, rhs_scale, scale_factor):
15651576
def dot_scaled(self, lhs: TensorTy, lhs_scale: TensorTy, lhs_format: str, rhs: TensorTy,
15661577
rhs_scale: Optional[TensorTy], rhs_format: str, acc: TensorTy | None, fast_math: bool,
15671578
lhs_k_pack: bool, rhs_k_pack: bool, out_dtype: tl.dtype) -> TensorTy:
1568-
assert lhs.type.is_block() and rhs.type.is_block()
1579+
assert lhs.type.is_block() and rhs.type.is_block(), "dot_scaled operands must be block tensors (not scalars)"
15691580
#TODO: validate types.
15701581
lhs_rank = len(lhs.shape)
15711582
rhs_rank = len(rhs.shape)
@@ -1604,7 +1615,11 @@ def dot_scaled(self, lhs: TensorTy, lhs_scale: TensorTy, lhs_format: str, rhs: T
16041615
acc_handle = self.builder.create_splat(ret_ty.to_ir(self.builder), _0)
16051616
else:
16061617
acc_handle = acc.handle
1607-
assert acc.type.shape == ret_ty.shape and acc.type.element_ty == out_dtype
1618+
assert acc.type.shape == ret_ty.shape, \
1619+
f"expected accumulator shape {ret_ty.shape}, got {acc.type.shape}"
1620+
assert acc.type.element_ty == out_dtype, \
1621+
f"expected accumulator dtype {out_dtype}, got {acc.type.element_ty}; " \
1622+
f"pass out_dtype={acc.type.element_ty} to use this accumulator dtype"
16081623
rhs_scale_handle = None if rhs_scale_is_none else rhs_scale.handle
16091624
lhs_scale_handle = None if lhs_scale_is_none else lhs_scale.handle
16101625

@@ -1842,7 +1857,7 @@ def make_tensor_descriptor(self, base: TensorTy, shape: List[TensorTy], strides:
18421857
raise ValueError(f"Expected {ndim} strides but got {len(strides)}")
18431858
if len(block_shape) != ndim:
18441859
raise ValueError(f"Expected block_shape to have {ndim} dimensions but got {len(strides)}")
1845-
assert isinstance(base.dtype, tl.pointer_type)
1860+
assert isinstance(base.dtype, tl.pointer_type), f"base must be a pointer type, got {base.dtype}"
18461861
elem_size = base.dtype.element_ty.primitive_bitwidth // 8
18471862
contig_dim_size = tl._unwrap_if_constexpr(block_shape[-1])
18481863
if contig_dim_size * elem_size < 16:
@@ -1860,7 +1875,7 @@ def make_tensor_descriptor(self, base: TensorTy, shape: List[TensorTy], strides:
18601875
# Check whether `block_shape` is static
18611876
block_shape = tl._unwrap_shape(block_shape)
18621877

1863-
assert isinstance(base.type, tl.pointer_type)
1878+
assert isinstance(base.type, tl.pointer_type), f"base must be a pointer type, got {base.type}"
18641879
type = tl.block_type(base.type.element_ty, block_shape)
18651880
base_handle = base.handle
18661881
is_signed_int = base.type.element_ty.is_int_signed()

0 commit comments

Comments
 (0)