[FPSAN] Preserve NaN payload bits in encoding

Author: jeffniu-openai

lib/Conversion/TritonInstrumentToLLVM/FpSanToLLVM.cpp

@@ -123,10 +123,14 @@ Value mixFloatToInt(ConversionPatternRewriter &rewriter, Location loc, Value u,
   PayloadMixConfig cfg = getPayloadMixConfig(floatTy);
   Value signFlip =
       selectUIntConstantOnSign(rewriter, loc, u, cfg.signMask, 0, cfg.signMask);
-  Value x = b.xor_(u, signFlip);
   Value mulA = createUIntConstant(rewriter, loc, u.getType(), cfg.mulA);
   Value magMask = createUIntConstant(rewriter, loc, u.getType(), cfg.magMask);
-  Value yMul = b.mul(x, mulA);
+  // Avoid patterns that InstCombine rewrites to `llvm.fabs`. LLVM specifies
+  // that `llvm.fabs` preserves the NaN quiet/signaling bit and payload, but
+  // NVPTX lowers it to PTX `abs.f32`, whose NaN result is unspecified. On
+  // Blackwell, `abs.f32` is observed to canonicalize signaling NaNs, corrupting
+  // FPSan payloads.
+  Value yMul = b.mul(u, mulA);
   Value y = b.and_(yMul, magMask);
   Value z = xorShiftRight(rewriter, loc, y, cfg.shift);
   Value mulB = selectUIntConstantOnSign(rewriter, loc, u, cfg.signMask,
@@ -142,11 +146,10 @@ Value unmixIntToFloat(ConversionPatternRewriter &rewriter, Location loc,
   PayloadMixConfig cfg = getPayloadMixConfig(floatTy);
   Value signFlip =
       selectUIntConstantOnSign(rewriter, loc, v, cfg.signMask, 0, cfg.signMask);
-  Value w = b.xor_(v, signFlip);
   Value magMask = createUIntConstant(rewriter, loc, v.getType(), cfg.magMask);
   Value mulBInv = selectUIntConstantOnSign(rewriter, loc, v, cfg.signMask,
                                            cfg.mulBPosInv, cfg.mulBNegInv);
-  Value zMul = b.mul(w, mulBInv);
+  Value zMul = b.mul(v, mulBInv);
   Value z = b.and_(zMul, magMask);
   Value y = inverseXorShiftRight(rewriter, loc, z, cfg);
   Value mulAInv = createUIntConstant(rewriter, loc, v.getType(), cfg.mulAInv);

python/test/gluon/test_fpsan.py

@@ -2703,6 +2703,43 @@ def loop_sum_kernel(x_ptr, out_ptr, N: tl.constexpr):
     _assert_payload_equal(reduce_out, loop_out)
 
 
+def test_f32_loop_preserves_snan_payload(device, fresh_knobs):
+    _require_cuda_backend(device)
+    if not is_cuda():
+        pytest.skip("regression is specific to NVPTX fabs lowering")
+
+    @triton.jit
+    def sum_kernel(x_ptr, out_ptr, BLOCK: tl.constexpr):
+        offsets = tl.arange(0, BLOCK)
+        acc = tl.zeros((BLOCK, ), tl.float32)
+        for i in range(3):
+            acc += tl.load(x_ptr + i * BLOCK + offsets)
+        tl.store(out_ptr + offsets, acc)
+
+    fresh_knobs.compilation.instrumentation_mode = "fpsan"
+    fresh_knobs.compilation.always_compile = True
+
+    block = 128
+    # The first two finite values sum to an sNaN; the zero row forces it through the next loop embed.
+    input_bits = np.zeros((3, block), dtype=np.int32)
+    input_bits[0].fill(0x1B0F577C)
+    input_bits[1].fill(0x65E031B7)
+    assert np.isfinite(input_bits.view(np.float32)).all()
+    x = torch.tensor(input_bits, dtype=torch.int32, device="cuda")
+    out = torch.empty((block, ), dtype=torch.int32, device="cuda")
+    sum_kernel[(1, )](
+        triton.TensorWrapper(x, dtype=torch.float32),
+        triton.TensorWrapper(out, dtype=torch.float32),
+        BLOCK=block,
+        num_warps=1,
+    )
+
+    expected = _expected_add_i32(input_bits[0], input_bits[1])
+    expected = _expected_add_i32(expected, input_bits[2])
+    assert np.all(_as_u32(expected) == np.uint32(0x7FA12345))
+    _assert_payload_equal(out, expected)
+
+
 @pytest.mark.skipif(not (is_hip_cdna3() or is_hip_cdna4()), reason="Requires CDNA3 or CDNA4")
 @pytest.mark.parametrize(("type_a", "type_b", "acc_type", "m", "n", "k", "instr_m", "instr_n", "instr_k", "k_width"),
                          _MFMA_DOT_CASES)

test/Conversion/tritoninstrument_to_llvm.mlir

@@ -115,8 +115,9 @@ tt.func private @experimental_gsan_tensordesc_info(
 module attributes {"ttg.num-warps" = 4 : i32, ttg.target = "cuda:90"} {
 // CHECK-LABEL: @experimental_fpsan_embed
 // CHECK-NOT: tti.experimental_fpsan_embed
-// CHECK: llvm.bitcast %arg0 : f32 to i32
-// CHECK: llvm.mul
+// CHECK: %[[RAW:.*]] = llvm.bitcast %arg0 : f32 to i32
+// CHECK-NOT: llvm.inline_asm
+// CHECK: llvm.mul %[[RAW]],
 // CHECK: llvm.xor
 tt.func private @experimental_fpsan_embed(%arg0: f32) -> i32 {
   %0 = tti.experimental_fpsan_embed %arg0 : (f32) -> i32
@@ -129,7 +130,7 @@ tt.func private @experimental_fpsan_embed(%arg0: f32) -> i32 {
 module attributes {"ttg.num-warps" = 4 : i32, ttg.target = "cuda:90"} {
 // CHECK-LABEL: @experimental_fpsan_unembed
 // CHECK-NOT: tti.experimental_fpsan_unembed
-// CHECK: llvm.mul
+// CHECK: llvm.mul %arg0,
 // CHECK: llvm.xor
 // CHECK: llvm.bitcast %{{.*}} : i32 to f32
 tt.func private @experimental_fpsan_unembed(%arg0: i32) -> f32 {