[CuTeDSL] Skip the numpy round-trip for in-range integer scalars - #3446
Open
wilyan09007 wants to merge 1 commit into
Open
[CuTeDSL] Skip the numpy round-trip for in-range integer scalars#3446wilyan09007 wants to merge 1 commit into
wilyan09007 wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #3443.
Integer.__init__sends every Python scalar throughint(np.array(x).astype(np_dtype)), and the float path callsnp.isnan/np.isinfon 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 theFloat32path, and its cost is unchanged here.Numbers
CPython 3.12.3, numpy 2.5.2,
nvidia-cutlass-dsl4.7.0, no GPU. Minimum of 5 runs of 100k constructions each:Int32(5)Int64(5)Int32(3.7)Boolean(1)Int32, a launch prologueInt32(2**32), wrapsThe 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
IntegerMetacomputes the range once per type at class creation, fromnp.iinfoon the dtype, and stores it as_exact_range. It deliberately follows the numpy dtype rather thanwidth, because the cast is what it has to agree with and the two do not always match:Booleanis declared signed at width 1, which putsmin/maxat -1/0 while the dtype isnp.bool_. It is special-cased to (0, 1), the only two valuesastype(np.bool_)leaves alone.WgmmaSmemDescandTcgen05SmemDescassignwidthafter the class is built, so theirmin/maxdescribe 64 bits while they still cast throughint32. Keying the check offmin/maxwould makeWgmmaSmemDesc(2**40)store 1099511627776 instead of the 0 the cast produces.minandmaxthemselves are untouched.Behavior
Unchanged, including the parts that are arguably surprising:
Int8(128)is -128,Uint8(-1)is 255), and theOverflowErrornumpy raises for integers too wide for any C type.Int4,Int128andUint128have 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.ValueErrorandOverflowErrorwith the same messages, except forBoolean, which has always folded them to 1 because it runsbool()beforeInteger.__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.pyis new, 14 tests, no GPU or compilation needed. The broadest one asserts the constructor againstint(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, theBooleanrange widened, and the NaN check dropped. Each is caught.test/python/CuTeDSL/test_struct_in_if.pypasses before and after (6 tests, realcute.compile).