[BENCH] Address and/or warnings in triton kernels (#6841)

Author: Jokeren

python/triton/compiler/code_generator.py

@@ -1300,8 +1300,15 @@ def visit_BoolOp(self, node: ast.BoolOp):
                 # expression so we do not append it to nontrivial_values.
             else:
                 if value.type.is_block():
-                    warnings.warn(
-                        "Logical operators 'and' and 'or' are deprecated for non-scalar tensors; please use '&' or '|' instead"
+                    lineno = getattr(node, "lineno", None)
+                    if lineno is not None:
+                        lineno += self.begin_line
+                    warnings.warn_explicit(
+                        "Logical operators 'and' and 'or' are deprecated for non-scalar tensors; please use '&' or '|' instead",
+                        category=UserWarning,
+                        filename=self.file_name,
+                        lineno=lineno,
+                        source=ast.unparse(node),
                     )
                 # not a constexpr so we must append it:
                 nontrivial_values.append(value)

python/triton_kernels/triton_kernels/matmul_ogs_details/_matmul_ogs.py

@@ -321,7 +321,7 @@ def _compute_writeback_idx(
     is_src_active = (src_idxs != -1).to(tl.int32)
     has_one_active = tl.sum(is_src_active, axis=1) == 1
 
-    need_finalize_scatter = mask_m and not has_one_active
+    need_finalize_scatter = mask_m & (~has_one_active)
     finalize_scatter_count = tl.sum(need_finalize_scatter.to(tl.int32))
     if finalize_scatter_count == 0:
         return

python/triton_kernels/triton_kernels/numerics_details/mxfp.py

@@ -155,13 +155,13 @@ def _downcast_to_mxfp(mx_tensor_ptr, stride_mxt_outer, stride_mxt_quant: tl.cons
 
     mask_src_quant = start_src_quant + offs_src_quant < quant_dim
     mask_n = start_out + offs_outer < outer_dim
-    full_mask_src = mask_src_quant and mask_n
+    full_mask_src = mask_src_quant & mask_n
 
     mask_mxt_quant = start_mx_quant + offs_mxt_quant < tl.cdiv(quant_dim, K_DIVISOR)
-    full_mask_mxt = mask_mxt_quant and mask_n
+    full_mask_mxt = mask_mxt_quant & mask_n
 
     scale_mask_k = start_mx_scale_quant + offs_scale_quant < tl.cdiv(quant_dim, 32)
-    full_scale_mask = scale_mask_k and mask_n
+    full_scale_mask = scale_mask_k & mask_n
 
     src_tensor_offsets = offs_src_quant * stride_src_quant + offs_outer * stride_src_outer
     mx_scale_offsets = offs_scale_quant * stride_mx_scale_quant + offs_outer * stride_mx_scale_outer
@@ -219,13 +219,13 @@ def _upcast_from_mxfp(out_ptr, stride_o_outer, stride_o_quant: tl.constexpr,
 
     mask_outer = start_out + offs_outer < outer_dim
     mask_out_quant = start_out_quant + offs_out_quant < quant_dim
-    full_mask_out = mask_out_quant and mask_outer
+    full_mask_out = mask_out_quant & mask_outer
 
     mask_src_quant = start_mxt_quant + offs_src_quant < tl.cdiv(quant_dim, K_DIVISOR)
-    full_mask_src = mask_src_quant and mask_outer
+    full_mask_src = mask_src_quant & mask_outer
 
     mask_scale = start_mx_scale_quant + offs_scale < tl.cdiv(quant_dim, 32)
-    full_scale_mask = mask_scale and mask_outer
+    full_scale_mask = mask_scale & mask_outer
 
     tensor_offsets = offs_src_quant * stride_tensor_quant + offs_outer * stride_tensor_outer
     scale_offsets = offs_scale * stride_scale_quant + offs_outer * stride_scale_outer

python/triton_kernels/triton_kernels/swiglu_details/_swiglu.py

@@ -66,7 +66,7 @@ def _swiglu(Out, OutExpectedScale, OutActualScale, OutChecksumScale, A, AScale,
             if pid_n * BLOCK_N + BLOCK_N <= N:
                 a_packed = tl.load(A + packed_offs, mask=mask_m[:, None], other=0.)
             else:
-                packed_mask = mask_m[:, None] and packed_mask_n[None, :]
+                packed_mask = mask_m[:, None] & packed_mask_n[None, :]
                 a_packed = tl.load(A + packed_offs, mask=packed_mask, other=0.)
         a_gelu, a_linear = tl.split(tl.reshape(a_packed, (BLOCK_M, BLOCK_N, 2)))
         # a gelu