[BUG] TensorSSA FP32->FP8 cvt doesn't work

[gau-nernst]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
A clear and concise description of what the bug is.

"""Reproduce CuTeDSL eliminating a two-element FP32-to-FP8 vector cast."""

import os

os.environ["CUTE_DSL_KEEP_PTX"] = "1"
os.environ["CUTE_DSL_DUMP_DIR"] = "./workspace/cutedsl_dump"
os.environ["CUTE_DSL_NO_CACHE"] = "1"
os.makedirs(os.environ["CUTE_DSL_DUMP_DIR"], exist_ok=True)

import cutlass
import cutlass.cute as cute
import torch
from cuda.bindings.driver import CUstream
from cutlass import Float8E4M3FN, Float32, Uint8


class CastKernel:
    def __init__(self, vectorized: bool):
        self.vectorized = vectorized

    @cute.jit
    def __call__(self, x: cute.Tensor, out: cute.Tensor, stream: CUstream):
        self.kernel(x, out).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)

    @cute.kernel
    def kernel(self, x: cute.Tensor, out: cute.Tensor):
        tid, _, _ = cute.arch.thread_idx()
        values = cute.make_rmem_tensor(2, Float32)
        values[0] = x[tid, 0] + Float32(1)
        values[1] = x[tid, 1] + Float32(2)

        quantized = cute.make_rmem_tensor(2, Float8E4M3FN)
        if cutlass.const_expr(self.vectorized):
            quantized.store(values.load().to(Float8E4M3FN))
        else:
            quantized[0] = values[0].to(Float8E4M3FN)
            quantized[1] = values[1].to(Float8E4M3FN)

        copy = cute.make_copy_atom(
            cute.nvgpu.CopyUniversalOp(), Float8E4M3FN, num_bits_per_copy=16
        )
        cute.copy(copy, quantized, out[tid, None])


class ScalarCastKernel(CastKernel):
    def __init__(self):
        super().__init__(False)


class VectorCastKernel(CastKernel):
    def __init__(self):
        super().__init__(True)


def compile_kernel(kernel_type):
    x = cute.runtime.make_fake_tensor(Float32, (32, 2), (2, 1), assumed_align=16)
    out = cute.runtime.make_fake_tensor(
        Float8E4M3FN, (32, 2), (2, 1), assumed_align=16
    )
    stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
    return cute.compile(
        kernel_type(), x, out, stream, options="--enable-tvm-ffi"
    )


def main():
    x = torch.arange(64, device="cuda", dtype=torch.float32).reshape(32, 2) / 8
    expected = (x + x.new_tensor([1, 2])).to(torch.float8_e4m3fn)
    outputs = []
    for kernel_type in (ScalarCastKernel, VectorCastKernel):
        out = torch.full_like(x, float("nan"), dtype=torch.float8_e4m3fn)
        compile_kernel(kernel_type)(x, out)
        outputs.append(out)
    torch.cuda.synchronize()

    torch.testing.assert_close(outputs[0].float(), expected.float(), rtol=0, atol=0)
    print("scalar conversion: PASS")
    try:
        torch.testing.assert_close(outputs[1].float(), expected.float(), rtol=0, atol=0)
    except AssertionError as error:
        print("vectorized conversion: FAIL (output remains NaN)")
        raise AssertionError(
            "scalar FP32->FP8 conversion passed; vectorized conversion failed"
        ) from error
    print("vectorized conversion: PASS")


if __name__ == "__main__":
    main()

Result

AssertionError: scalar FP32->FP8 conversion passed; vectorized conversion failed

Inspecting the generated PTX

//
// Generated by NVIDIA NVVM Compiler
//
// Compiler Build ID: CL-36699951
// Cuda compilation tools, release 13.1, V13.1.66
// Based on NVVM 21.0.0
//

.version 9.1
.target sm_103a
.address_size 64

	// .globl	kernel_cutlass_kernel___main__VectorCastKernel_object_at__tensorptrf32gmemalign16o32221_tensorptrf8E4M3FNgmemalign16o32221_0

.visible .entry kernel_cutlass_kernel___main__VectorCastKernel_object_at__tensorptrf32gmemalign16o32221_tensorptrf8E4M3FNgmemalign16o32221_0(
	.param .align 8 .b8 kernel_cutlass_kernel___main__VectorCastKernel_object_at__tensorptrf32gmemalign16o32221_tensorptrf8E4M3FNgmemalign16o32221_0_param_0[8],
	.param .align 8 .b8 kernel_cutlass_kernel___main__VectorCastKernel_object_at__tensorptrf32gmemalign16o32221_tensorptrf8E4M3FNgmemalign16o32221_0_param_1[8]
)
.reqntid 32, 1, 1
{


	ret;

}

The compiler eliminates the entire kernel body

Steps/Code to reproduce bug
Follow this guide http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports to craft a minimal bug report. This helps us reproduce the issue you're having and resolve the issue more quickly.

Expected behavior
A clear and concise description of what you expected to happen.

Environment details (please complete the following information):

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]

CuteDSL 4.6.0, GB300

Additional context
Add any other context about the problem here.

[brandon-yujie-sun]

@gau-nernst thanks for reporting the issue. We will take a look