@@ -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