Skip to content

[CuTeDSL] Skip the numpy round-trip for in-range integer scalars - #3446

Open
wilyan09007 wants to merge 1 commit into
NVIDIA:mainfrom
wilyan09007:fix/issue-3443
Open

[CuTeDSL] Skip the numpy round-trip for in-range integer scalars#3446
wilyan09007 wants to merge 1 commit into
NVIDIA:mainfrom
wilyan09007:fix/issue-3443

Conversation

@wilyan09007

Copy link
Copy Markdown

Fixes #3443.

Integer.__init__ sends every Python scalar through int(np.array(x).astype(np_dtype)), and the float path calls np.isnan/np.isinf on a Python float before that. The cast is there to reproduce C wrap-around, but it is an identity whenever the value already fits the target dtype, which is the usual case for kernel arguments. This range checks first and keeps the cast for the values that really do wrap.

Float.__init__ already dropped its numpy conversion (the replaced code is still commented in place at the top of that method), so this is the same treatment applied to the integer side. There is nothing equivalent left to remove on the Float32 path, and its cost is unchanged here.

Numbers

CPython 3.12.3, numpy 2.5.2, nvidia-cutlass-dsl 4.7.0, no GPU. Minimum of 5 runs of 100k constructions each:

construction before after
Int32(5) 1165 ns 607 ns 1.9x
Int64(5) 1056 ns 651 ns 1.6x
Int32(3.7) 2450 ns 693 ns 3.5x
Boolean(1) 1281 ns 805 ns 1.6x
8x Int32, a launch prologue 10128 ns 4978 ns 2.0x
Int32(2**32), wraps 1170 ns 1268 ns 0.9x

The last row is the tradeoff: a value that wraps now pays one extra comparison before reaching the cast, roughly 100 ns against the 1.2 us the cast already costs.

Where the bounds come from

IntegerMeta computes the range once per type at class creation, from np.iinfo on the dtype, and stores it as _exact_range. It deliberately follows the numpy dtype rather than width, because the cast is what it has to agree with and the two do not always match:

  • Boolean is declared signed at width 1, which puts min/max at -1/0 while the dtype is np.bool_. It is special-cased to (0, 1), the only two values astype(np.bool_) leaves alone.
  • WgmmaSmemDesc and Tcgen05SmemDesc assign width after the class is built, so their min/max describe 64 bits while they still cast through int32. Keying the check off min/max would make WgmmaSmemDesc(2**40) store 1099511627776 instead of the 0 the cast produces.

min and max themselves are untouched.

Behavior

Unchanged, including the parts that are arguably surprising:

  • Out-of-range values keep numpy's wrap-around (Int8(128) is -128, Uint8(-1) is 255), and the OverflowError numpy raises for integers too wide for any C type.
  • Int4, Int128 and Uint128 have no numpy dtype and still trip the same assert. The shortcut sits behind that guard rather than in front of it, so they do not quietly gain support.
  • NaN and infinity still raise ValueError and OverflowError with the same messages, except for Boolean, which has always folded them to 1 because it runs bool() before Integer.__init__ sees the value.

Checked by snapshotting every integer and float type against 57 inputs (boundaries, wrap-around candidates, bools, floats, NaN and both infinities) plus construction from other DSL scalars, recording either the stored value or the exact exception, and diffing the run before the change against the run after: 1505 cases, no differences.

Tests

test/python/CuTeDSL/test_numeric_construction.py is new, 14 tests, no GPU or compilation needed. The broadest one asserts the constructor against int(np.array(x).astype(dtype)) directly, so the claim this PR rests on is the thing under test. The rest pin the wrap-around, the truncation direction, the NaN and infinity handling, the types with no numpy dtype, and the dtype-versus-width case above.

I also ran them against four deliberately broken variants to confirm they fail: bounds keyed off min/max, the range guard removed, the Boolean range widened, and the NaN check dropped. Each is caught.

test/python/CuTeDSL/test_struct_in_if.py passes before and after (6 tests, real cute.compile).

Integer.__init__ converted every Python scalar with
int(np.array(x).astype(np_dtype)). That cast is there to reproduce C
wrap-around, but it costs more than the rest of the constructor put
together, and it is an identity whenever the value already fits the
target dtype. Range check first and fall back to numpy only for the
values that really do wrap.

The NaN and infinity guard on the float path called np.isnan/np.isinf on
Python floats. Each is a ufunc dispatch roughly an order of magnitude
slower than the math module equivalent, and Int32(3.7) paid both on top
of the cast.

Bounds are taken from np.iinfo once per type at class creation and kept
in IntegerMeta._exact_range. They follow the numpy dtype rather than
width, because the cast is what they have to agree with: Boolean is
declared signed at width 1, and WgmmaSmemDesc widens width after the
class is built while keeping an int32 dtype.

Measured on CPython 3.12 with numpy 2.5:

  Int32(5)                     1165 ns -> 607 ns
  Int64(5)                     1056 ns -> 651 ns
  Int32(3.7)                   2450 ns -> 693 ns
  Boolean(1)                   1281 ns -> 805 ns
  8x Int32, launch prologue   10128 ns -> 4978 ns

Values that wrap now pay one extra comparison, about 100 ns against the
1.2 us the cast already costs.

Signed-off-by: William <wilyan090@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant