[BUG][CuTe DSL] Constructing cutlass Numeric scalars from Python values costs ~0.3-0.8us each

[reubenconducts]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
Calling cutlass.Int32(val) costs ~750ns and cutlass.Float32(val) costs ~270ns. This is quite noticeable, on the order of kernel launch latency itself. Integer.__init__ (cutlass/base_dsl/typing.py) converts a plain Python int through a numpy round-trip: x_val = int(np.array(x).astype(np_dtype)) which alone costs ~400us.

Steps/Code to reproduce bug

import timeit

import numpy as np

import cutlass

INT32_MIN, INT32_MAX = -(2**31), 2**31 - 1


def numpy_path(x: int) -> int:
    # The conversion Integer.__init__ performs today.
    return int(np.array(x).astype(np.int32))


def fast_path(x: int) -> int:
    # Range check first; fall back to numpy only for wrap-around.
    if INT32_MIN <= x <= INT32_MAX:
        return int(x)  # int() normalizes bool -> int, matching the numpy path
    return numpy_path(x)


def check_equivalence():
    edge_cases = [
        0, 1, -1, 5, True, False,
        INT32_MIN, INT32_MAX, INT32_MIN - 1, INT32_MAX + 1,
        2**31, 2**32, -(2**32), 2**63 - 1, 123456789,
    ]
    for x in edge_cases:
        assert fast_path(x) == numpy_path(x), x
    print(f"fast path matches numpy path on {len(edge_cases)} edge cases")


def bench(label, fn, n=100_000):
    t = min(timeit.repeat(fn, number=n, repeat=5)) / n
    print(f"  {label:40s} {t * 1e9:7.0f} ns")


def main():
    check_equivalence()
    print("per-construction cost:")
    bench("cutlass.Int32(5)", lambda: cutlass.Int32(5))
    bench("cutlass.Int64(5)", lambda: cutlass.Int64(5))
    bench("cutlass.Float32(0.5)", lambda: cutlass.Float32(0.5))
    bench("int(np.array(5).astype(np.int32))", lambda: numpy_path(5))
    bench("range-checked fast path", lambda: fast_path(5))

    # Context: a realistic launch prologue wrapping 8 scalar kernel arguments.
    scalars = [1, 8, 128, 576, 4, 32, 1024, 7]
    bench(
        "8x Int32 (typical launch prologue)",
        lambda: [cutlass.Int32(s) for s in scalars],
        n=20_000,
    )


if __name__ == "__main__":
    main()

Expected behavior
These casts shouldn't take this long. A workaround is to define, e.g.,

@lru_cache
def _int32(val: int) -> cutlass.Int32:
	return cutlass.Int32(val)

Environment details (please complete the following information):

• nvidia-cutlass-dsl 4.6.0.dev0 and 4.7.0, no GPU needed.