[FRONTEND] fix semantic issue for max_num_imprecise_acc (triton-lang#2835)

Author: ptillet

python/test/unit/language/test_core.py

@@ -2567,6 +2567,31 @@ def kernel(X, stride_xm, stride_xk, Y, stride_yk, stride_yn, W, stride_wn, strid
                 'mma.sync.aligned.m16n8k32.row.col.satfinite.s32.s8.s8.s32' in ptx
 
 
+def test_max_num_imprecise_acc(device):
+    capability = torch.cuda.get_device_capability()
+    if capability != (9, 0):
+        return
+
+    @triton.jit
+    def kernel(X, Y, Z, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+               MAX_NUM_IMPRECISE_ACC: tl.constexpr):
+        off_m = tl.arange(0, BLOCK_M)
+        off_n = tl.arange(0, BLOCK_N)
+        off_k = tl.arange(0, BLOCK_K)
+        x = tl.load(X + off_m[:, None] * BLOCK_K + off_k[None, :])
+        y = tl.load(Y + off_k[:, None] * BLOCK_N + off_n[None, :])
+        z = tl.load(Z + off_m[:, None] * BLOCK_N + off_n[None, :])
+        z = tl.dot(x, y, acc=z, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC)
+        tl.store(Z + off_m[:, None] * BLOCK_N + off_n[None, :], z)
+
+    M, N, K, num_warps, MAX_NUM_IMPRECISE_ACC = 128, 128, 128, 4, 64
+    x = torch.zeros((M, K), dtype=torch.float8_e5m2, device=device)
+    y = torch.zeros((K, N), dtype=torch.float8_e5m2, device=device)
+    z = torch.zeros((M, N), dtype=torch.float32, device=device)
+    h = kernel[(1, 1)](x, y, z, M, N, K, MAX_NUM_IMPRECISE_ACC, num_warps=num_warps)
+    assert h.asm["ptx"].count("add.f32") == (M * N) // (32 * num_warps) * (K / MAX_NUM_IMPRECISE_ACC)
+
+
 @pytest.mark.parametrize('in_dtype', ['float32'])
 def test_dot_mulbroadcastred(in_dtype, device):
     capability = torch.cuda.get_device_capability()

python/triton/language/semantic.py

@@ -1306,10 +1306,11 @@ def assert_dtypes_valid(lhs_dtype, rhs_dtype, options):
         assert acc.type == ret_ty
 
     # max_num_imprecise_acc only applies to fp8 -> fp32 dot on sm_90
-    if lhs.dtype.is_fp8() and rhs.dtype.is_fp8() and max_num_imprecise_acc is None:
-        max_num_imprecise_acc = builder.options.max_num_imprecise_acc_default
-    else:
-        max_num_imprecise_acc = 0
+    if max_num_imprecise_acc is None:
+        if lhs.dtype.is_fp8() and rhs.dtype.is_fp8():
+            max_num_imprecise_acc = builder.options.max_num_imprecise_acc_default
+        else:
+            max_num_imprecise_acc = 0
 
     return tl.tensor(builder.create_dot(lhs.handle, rhs.handle, acc_handle, allow_tf32, max_num_imprecise_acc), ret_ty)