feat: ATen-style ops dispatch, compiler, and full test suite#5
Open
booth-algo wants to merge 2 commits intokev/aten-testbenchfrom
Open
feat: ATen-style ops dispatch, compiler, and full test suite#5booth-algo wants to merge 2 commits intokev/aten-testbenchfrom
booth-algo wants to merge 2 commits intokev/aten-testbenchfrom
Conversation
This was referenced Mar 31, 2026
booth-algo
force-pushed
the
kev/aten-testbench
branch
from
bc2ff7b to
d774500
Compare
booth-algo
force-pushed
the
kev/aten-ops
branch
from
7b29c38 to
7d66e7e
Compare
…ler, SubMatrixManager) 3-layer compilation stack for generating PLENA ISA from high-level ops: - PLENAProgram: tensor proxy API, HBM auto-allocation, scoped naming - DeveloperCompiler: register allocation, ISA string emission, ASM templates - SubMatrixManager: VRAM/MRAM/HBM memory layout, sub-block addressing Supporting infrastructure: - symbol_table, config_utils, emulator_runner, check_mem - Rust emulator: new opcodes (V_SHFT_V, H_STORE_V_PART) - justfile: test recipes for all operator tests - CLAUDE.md: project context + CI check instructions Co-Authored-By: Ziqian Gao <zg1223@ic.ac.uk>
ATen operator dispatch system (plena/ops/): - OpRegistry with CPU (golden reference) and PLENA (ISA generation) backends - native_ops.yaml declarative op definitions - Registered ops: softmax, linear, rms_norm, layer_norm, ffn, flash_attention, conv2d, embedding_add, rope ATen compiler (plena/compiler/aten_compiler.py): - Traces nn.Module via torch.export, walks ATen graph - FFN fusion detection, residual save/restore pre-pass - Dispatches to PLENA ops backends to produce ISA code Test suite (18 test recipes, all passing): - Primitive op tests, ATen compiler e2e tests - Real-model: SmolLM2-135M, CLM-60M, LLaDA-8B, SmolVLM2 - Model layer test builder with HF loading + MXFP8 quantization
booth-algo
force-pushed
the
kev/aten-ops
branch
from
7d66e7e to
050220c
Compare
booth-algo
force-pushed
the
kev/aten-testbench
branch
from
d774500 to
4f38fcb
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.
Summary
Depends on: #4 (kev/aten-testbench)
ATen Operator Dispatch (
plena/ops/)OpRegistrywith CPU (golden reference) and PLENA (ISA generation) backendsnative_ops.yamldeclarative op definitionssoftmax·linear·rms_norm·layer_norm·ffn·flash_attention·conv2d·embedding_add·ropeATen Compiler (
plena/compiler/aten_compiler.py)nn.Moduleviatorch.export, walks ATen graphFull Test Suite (18 recipes, all passing)
Real-Model Tests
Test plan