feat(codegen,runtime): emit the unary kinds the DSL already exposes - #72
Open
bigSheep123 wants to merge 1 commit into
Open
feat(codegen,runtime): emit the unary kinds the DSL already exposes#72bigSheep123 wants to merge 1 commit into
bigSheep123 wants to merge 1 commit into
Conversation
`tf.exp`, `tf.exp2`, `tf.log`, `tf.log2`, `tf.abs`, `tf.ceil` and `tf.round` are all reachable from the DSL and all type-infer and evaluate, but `_UNARY_TAG` carried rows for only `RSQRT`, `NEG`, `RELU` and `SQUARE`. Everything else failed at emission with a bare `KeyError: <UnaryKind.LOG: 'log'>` -- writable and checkable, but not compilable. Add the seven missing functors and their tag rows, following the existing pattern. Two choices worth stating: - The transcendentals use the precise `<math.h>` entries (`expf`, `logf`, ...) rather than the `__`-prefixed intrinsics, so a compiled kernel agrees with the evaluator's torch oracle instead of being a little faster and a little different. A kernel that would rather have the intrinsic can still reach for it explicitly. - `round` uses `rintf`, which breaks halfway cases to even, matching torch. `roundf` rounds them away from zero and would disagree at every `.5`; the test input lands exactly on those boundaries so the distinction is covered rather than assumed. `CAST` and `NOT` stay absent on purpose, and the table now says so: `CAST` has its own handler, and `NOT` is a boolean operation rather than the float path these functors assume. Each test kernel is a separate top-level function because the compiled artifact is cached per function -- a factory emitting several same-named kernels hands every op whichever one compiled first. pytest tests/ -q: 644 passed before, 651 after (+7 new, no regressions). Verified failing on main first: all seven raise `KeyError` at emission.
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.
feat(codegen,runtime): emit the unary kinds the DSL already exposes
Fixes #71.
What was wrong
tf.exp,tf.exp2,tf.log,tf.log2,tf.abs,tf.ceilandtf.roundareall reachable from the DSL and all type-infer and evaluate, but
_UNARY_TAGcarried rows for only
RSQRT,NEG,RELUandSQUARE. Everything elsefailed at emission with a bare
KeyError: <UnaryKind.LOG: 'log'>— writable andcheckable, but not compilable.
The change
The seven missing runtime functors and their tag rows, following the existing
pattern in
unary.cuh.Two choices worth calling out:
<math.h>entries (expf,logf,exp2f,log2f) rather than the__-prefixed intrinsics, so a compiledkernel agrees with the evaluator's torch oracle instead of being slightly
faster and slightly different. A kernel that would rather have the intrinsic
can still reach for it explicitly.
roundusesrintf, which breaks halfway cases to even, matchingtorch.round.roundfrounds them away from zero and would disagree at every.5.CASTandNOTstay absent deliberately, and the table now says so rather thanleaving them looking pending:
CASThas its own handler, andNOTis a booleanoperation rather than the float path these functors assume.
Tests
tests/integration/test_unary_codegen.pycompiles one kernel per op andcompares against the torch oracle on GPU. Verified to fail on
mainfirst —all seven raise
KeyErrorat emission — and pass after.The inputs are chosen per op rather than shared:
log/log2get a strictlypositive domain, and
roundgets a grid landing exactly on.5boundaries sothe rounding-mode choice above is actually exercised instead of assumed.
Each kernel is a separate top-level function because the compiled artifact is
cached per function — a factory emitting several same-named kernels hands every
op whichever one compiled first.
Verification
Every
.pre-commit-config.yamlhook clean: ruff, spec-rules, spec-refs,spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all C++ files.