operators: propagate dtype into AIERuntimeArgSpec - #179
Open
atassis wants to merge 1 commit into
Open
Conversation
strided_copy/repeat/gemm build AIERuntimeArgSpec without a dtype kwarg, so every returned spec silently reports the dataclass default (bfloat16) regardless of the operator's own dtype field(s). The two real consumers (test_utils.py's XRTTensor(spec.shape, dtype=spec.dtype), sequence.py's calculate_buffer_layout) size buffers off spec.dtype, so a non-default-dtype caller under- or over-allocates by the itemsize ratio -- GEMM(dtype_in='i8', dtype_out='i32') declared a 2-byte C buffer for an actual 4-byte int32 output. Latent: every in-tree caller uses the default today. Mirrors dequant/op.py's existing correct pattern.
atassis
force-pushed
the
fix/argspec-dtype-not-propagated
branch
from
August 31, 2026 21:36
2db0521 to
a8b0c95
Compare
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
1 participant
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.
Problem
AIERuntimeArgSpec.dtypedefaults tobfloat16, and none of the threeget_arg_spec()implementations passdtype=. So the spec reports bf16 regardless of what the design actually moves.Concretely,
GEMM(dtype_in='i8', dtype_out='i32')reports its C arg spec as bfloat16 -- 32768 bytes -- while the design DMAs int32, 65536 bytes. A caller sizing a host buffer from the spec under-allocates by 2x. The same 2x applies toStridedCopyandRepeatwith a non-default dtype.Two in-tree consumers size buffers off that value:
test_utils.py'sXRTTensorsizing andsequence.py'scalculate_buffer_layout.The defect is latent today -- a grep of every constructor call site in
iron/shows nothing currently passes a non-default dtype -- so this is a landmine rather than a live failure.Fix
Pass the dtype through at the 7 call sites across 3 operator files:
dtype=self.dtypeforstrided_copyandrepeat,dtype=str_to_dtype(self.dtype_in/out)forgemm.The two supporting commits are what the new test needs:
comparison.pygains a strict comparison so exact equality is expressible, andconftest.pyno longer aborts the whole session on an unparametrized test.Test
New
iron/tests/common/arg_spec_dtype.py, 6 device-free cases.Verified red-then-green by hand rather than only after: 5 of 6 fail before the change, 6 of 6 pass after. (Run by importing the module and calling the test functions directly --
conftest.pyresolves the device at collection time, so evenpytest --collect-onlywould open the NPU.)Both consumers confirmed to compute correct byte counts after the fix, and the change is a byte-for-byte no-op on every in-tree caller's default-dtype path.
blackclean.Known limitation
pr/strided-copy-transfer-sizealso rewritesget_arg_spec()for an unrelated shape bug. Whichever lands second will need a small reconcile.