Runtime TXN generation library and compiler support#3002
Runtime TXN generation library and compiler support#3002
Conversation
Introduce new operations that accept SSA values instead of static attributes, enabling runtime parameterization of NPU sequences: - aiex.npu.dyn_write32: Dynamic write with SSA address and value - aiex.npu.dyn_maskwrite32: Dynamic masked write with SSA operands - aiex.npu.dyn_dma_memcpy_nd: Fully dynamic N-D DMA with SSA sizes/strides - aiex.npu.dyn_sync: Dynamic synchronization with SSA tile/channel These operations can be lowered to templated C++ code for runtime transaction generation, allowing a single compiled artifact to support multiple problem sizes determined at runtime. Added verification to ensure SSA operands have correct types (index or signless integers). Maximum 4 dimensions enforced for DMA operations to match hardware constraints. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Implements the dynamic runtime sequences infrastructure: Phase 1: Extract TXN instruction encoding from AIETargetNPU.cpp into include/aie/Runtime/TxnEncoding.h - a header-only C++ library with zero MLIR/LLVM dependencies. Refactor AIETargetNPU.cpp to use it. Phase 2: Add ConvertAIEXToEmitC pass that lowers AIEX runtime sequence ops (both static npu.write32/blockwrite/sync/address_patch and dynamic npu.dyn_write32/dyn_maskwrite32/dyn_sync) plus SCF/arith ops into EmitC dialect. The EmitC IR calls aie_runtime::txn_append_* functions. Phase 3: Wire aie-translate --aie-generate-txn-cpp translation that runs the NPU lowering pipeline then the EmitC pass, producing compilable C++ that generates TXN binaries at runtime. Phase 4: Add test_dynamic.cpp example that uses the generated C++ instead of loading insts.bin. Verified bit-for-bit identical TXN output and PASS on NPU hardware. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Extract the NPU lowering pipeline (7 passes) into a shared populateNpuLoweringPipeline() function used by both aiecc and aie-translate, eliminating the duplicated pass list. Unify the host test executable: test_dynamic.cpp is deleted and test.cpp gains a USE_DYNAMIC_TXN compile flag that sets a generate_instr callback on the args struct in xrt_test_wrapper.h. Both static (insts.bin) and dynamic (generated TXN) paths now share the same XRT setup, buffer management, verification, and timing infrastructure. Add --aie-generate-txn-cpp and --txn-cpp-name flags to aiecc so C++ TXN generation is accessible through the compiler driver alongside --aie-generate-npu-insts. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Add test_dynamic_size.mlir using aiex.npu.dyn_write32 with BLOCKWRITE for BD configuration (required for address_patch compatibility) and dynamic buffer_length as a runtime parameter. Add --dynamic-size flag to test executable. When set, the host passes the transfer size to generate_txn_sequence() at runtime instead of using a pre-compiled instruction binary. The core loops processing fixed-size ObjectFIFO tiles, so any multiple of the tile size works. Demonstrated: XCLBIN compiled once at 1024-byte tile size, single executable runs correctly at 1024, 2048, 3072, and 4096 bytes with zero recompilation. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Extend NpuWrite32Op, NpuMaskWrite32Op, and NpuSyncOp with optional SSA operands (dyn_address, dyn_value, etc.) so a single op handles both compile-time constant and runtime-parameterized forms. Delete the 4 separate Dyn ops (NpuDynWrite32Op, NpuDynMaskWrite32Op, NpuDynSyncOp, NpuDynDmaMemcpyNdOp) that previously duplicated them. - Add AttrSizedOperandSegments trait and custom parse/print/verify - Add custom builders preserving existing call-site signatures - Merge EmitC conversion handlers (static vs dynamic dispatch) - Add error guards in NPU binary translation for dynamic operands - Add Python wrappers: npu_write32_dynamic, npu_maskwrite32_dynamic, npu_sync_dynamic - Replace hand-written test_dynamic_size.mlir with Python design file (passthrough_kernel_dynamic.py) 100% backward compatible: all 294 existing LIT tests pass unchanged. Co-Authored-By: Claude Opus 4.6 <[email protected]>
…amic Worker Extend IRON APIs to support dynamic (RTP-based) loop bounds: - Worker: add dynamic_objfifo_lowering parameter - Runtime: extend sequence() for mixed array/scalar types, add write_rtp() - RuntimeScalar: new class for scalar runtime sequence parameters - RtpWriteTask: new task class wrapping npu_rtp_write - single_core_iron_dynamic.py: IRON-level dynamic GEMM example Co-Authored-By: Claude Opus 4.6 <[email protected]>
… and LIT tests Extract the 2 dynamic designs (low-level, IRON) and their test harness from single_core/ into a new single_core_dynamic/ directory. Add a new placed dynamic variant using shim_dma_single_bd_task. Create LIT tests for all 3 variants and a passthrough_kernel dynamic LIT test. Fix dynamic_gemm_txn.h: add missing MASKWRITE before S2MM queue push (required for XRT completion token) and restructure C output BDs to use the batched pingpong pattern matching the static compiler. Verified on NPU2 hardware: all 3 dynamic sizes (32x32x32, 64x64x64, 128x128x128) PASS with a single XCLBIN for both low-level and placed variants. Co-Authored-By: Claude Opus 4.6 <[email protected]>
The dynamic TXN generator now parses the static instruction stream (generated alongside the XCLBIN) to discover the RTP buffer address and S2MM control register values. This makes the dynamic test harness work with all 3 design variants (low-level, placed, IRON) since each may place the RTP buffer at a different address. Previously the IRON variant failed because its buffer allocator placed the RTP at 0x204d00 while the code hardcoded 0x200600. All 3 variants now PASS on NPU2 at 32x32x32, 64x64x64, 128x128x128. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Co-Authored-By: Claude Opus 4.6 <[email protected]>
Enable auto-generated C++ TXN code from MLIR runtime sequences with SSA parameters, allowing a single XCLBIN to run matrix multiplications at any M/K/N (multiples of 32) determined at runtime. Key changes: - Add IsolatedFromAbove trait to RuntimeSequenceOp, preventing SCF-to-CF from entering runtime sequences while still lowering core bodies. Also prevents constant hoisting across the isolation boundary. - Extend DmaToNpuPattern with dynamic code path: when sizes/strides are SSA values, compute BD words via arith ops and emit npu.write32. Fixes bf16 d0_size (multiply-first-then-divide), stride underflow guards (size>1 check), and repeat_count off-by-one. - Extend EmitC conversion for scf.for with iter_args (VariableOp + LoadOp + AssignOp pattern), scf.if with results, and new arith ops (TruncI, ExtUI, ExtSI, MinSI, MaxSI). Add pre-scan for values hoisted outside runtime_sequence and cross-reference fixup pass. - Add dyn_arg_plus to NpuAddressPatchOp and dyn_value to NpuWriteRTPOp for runtime-parameterized buffer offsets and RTP writes. - Scope SCF-to-CF in aiecc via markOpRecursivelyLegal on RuntimeSequenceOp, and disable cross-region constant CSE in AIEVectorTransferLoweringPass. - Add unified aiecc compilation (--aie-generate-xclbin + --aie-generate-txn-cpp) producing both XCLBIN and C++ TXN from the same MLIR with identical buffer addresses. Verified on NPU Strix Halo: 32x32x32, 64x64x64, 64x32x64, 96x96x96, 128x64x128, 128x128x128 all PASS against reference matmul. Co-Authored-By: Claude Opus 4.6 <[email protected]>
Adapt to upstream API changes and fix IsolatedFromAbove interaction with MLIR's dialect conversion infrastructure: - Move convert-vector-to-aievec from resource allocation pipeline to per-core LLVM lowering, preventing vectorization of scalar arith ops (e.g. arith.minsi → aievec.min) inside runtime_sequence - Walk RuntimeSequenceOps explicitly in DmaToNpu, DMATasksToNPU, LowerSetLock, SubstituteShimDMA, since applyPartialConversion no longer descends into IsolatedFromAbove regions in newer LLVM - Skip materialize pass in AIETranslateToCppTxn (runtime_sequence is already in final form) - Add type casts in EmitC yield handler for mixed i32/opaque types - Disable constant CSE in AIEMaterializeBDChains - Update Python API: link_with on external_func, TraceShimRouting enum - Add hasVerifier to RunOp, dyn_arg_plus to AIEInsertTraceFlows All 6 GEMM sizes verified on NPU Strix Halo after rebase. Co-Authored-By: Claude Opus 4.6 <[email protected]>
- Fix getAsValue to zero-extend narrow values (was truncate-only) - Fix yieldTargets ArrayRef invalidation by using stack directly - Fail on unsupported ops in EmitC instead of silently emitting comments - Fix 0x80000000u token bit cast to int32_t - Remove dead preSCFModule global variable - Add IsolatedFromAbove negative test for RuntimeSequenceOp - Add bf16 d0_stride hardware constraint comment - Remove dead vectorized variable, name ROWS_PER_BLOCK constant - Deduplicate trace event list into module-level constant - Fix npu_time_min initialization to numeric_limits::max Co-Authored-By: Claude Opus 4.6 <[email protected]>
Stale submodule pointer from before rebase. Co-Authored-By: Claude Opus 4.6 <[email protected]>
AIETargetCppTxn.cpp and AIENpuLowering.cpp link AIEXTransforms, which uses BdIdGenerator from AIETransforms. Without this transitive dependency, static Release builds fail with undefined references to BdIdGenerator::nextBdId etc. in AIEAssignRuntimeSequenceBDIDs.cpp. Co-Authored-By: Claude Opus 4.6 <[email protected]>
2c839d4 to
93acc43
Compare
IsolatedFromAbove broke 62 existing tests that reference device-scope values (tiles, locks) from inside runtime_sequence. Instead, protect against constant hoisting by stripping runtime_sequences from LLVM lowering clones (where convert-vector-to-aievec's canonicalizer was the source of the hoisting). The markOpRecursivelyLegal SCF→CF scoping and enableConstantCSE(false) in AIEVectorTransferLowering remain as the primary guards. Co-Authored-By: Claude Opus 4.6 <[email protected]>
47bb7b4 to
c97ce07
Compare
Co-Authored-By: Claude Opus 4.6 <[email protected]>
3e32ea1 to
5286ec4
Compare
Coverage ReportCreated: 2026-03-26 23:50Click here for information about interpreting this report.
Generated by llvm-cov -- llvm version 18.1.3 |
Co-Authored-By: Claude Opus 4.6 <[email protected]>
53af3d9 to
934d90b
Compare
| f.write( | ||
| """ |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| f.write( | |
| """ | |
| f.write(""" |
| """ | ||
| ) |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| """ | |
| ) | |
| """) |
| f.write( | ||
| """extern "C" { |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| f.write( | |
| """extern "C" { | |
| f.write("""extern "C" { |
| }""" | ||
| ) |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| }""" | |
| ) | |
| }""") |
| f.write( | ||
| """extern "C" { |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| f.write( | |
| """extern "C" { | |
| f.write("""extern "C" { |
| }""" | ||
| ) |
There was a problem hiding this comment.
[black] reported by reviewdog 🐶
| }""" | |
| ) | |
| }""") |
Co-Authored-By: Claude Opus 4.6 <[email protected]>
There was a problem hiding this comment.
Nice to see this starting to take shape. The biggest question is if we want to deprecate the attributes; I'd be in favor of it, although it would mean touching potentially a lot of tests (but the actual code would be smaller).
The GEMM test unfortunately doesn't seem to use the added infrastructure. I think we are already on the same page, but just to make sure, what I'm envisioning looks more like this:
User writes a single runtime sequence in MLIR/Python (pseudocode):
aiex.runtime_sequence @my_sequence(%A: memref, %B: memref, %C: memref, %param_M: int, %param_K: int, %param_N: int) {\
...
aie.dma_memcpy_nd(...)
...
}
User calls compiler roughly like so:
aie-opt --aie-to-cpp aie.mlir -o my_runtime_sequence.h
which produces something like this using emitC (important -- this is compiler generated from the above MLIR, not manually written like in the GEMM test):
#include <txn_encoding.h>
std::vector<uint32_t> my_sequence(void *A, void *B, void *C, int param_M, int param_K, int param_N) {
std::vector<uint32_t> txn;
aie_runtime::txn_append_write32(txn, param_M, ...)
...
}
and then can use that generated file in their test.cpp like so:
#include "my_runtime_sequence.h"
int main(){
// setup XRT
xrt::kernel my_kernel = // get out of xclbin
std::vector<uint32_t> insts = my_sequence( my params ... )
my_kernel(insts, a, b, c);
}
So ideally the GEMM test's test.cpp at the end of this would not look significantly more complicated than the existing ones do.
Again, cool to see this taking shape!
| // See https://llvm.org/LICENSE.txt for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
| // | ||
| // (c) Copyright 2025 Advanced Micro Devices, Inc. |
| OptionalAttr<I32Attr>:$value, | ||
| Optional<I32>:$dyn_value | ||
| ); | ||
| let results = (outs ); | ||
| let assemblyFormat = [{ `(` $buffer `,` $index `,` $value `)` attr-dict | ||
| }]; | ||
| let hasCustomAssemblyFormat = 1; | ||
| let hasVerifier = 1; | ||
| let description = [{ | ||
| rtp write operator | ||
| rtp write operator. | ||
| When `dyn_value` is provided, it supplies the RTP value at runtime | ||
| instead of the static `value` attribute. | ||
| }]; | ||
| let extraClassDeclaration = [{ | ||
| bool hasDynamicValue() { return getDynValue() != nullptr; } | ||
| }]; |
There was a problem hiding this comment.
I'm a little worried about code bloat with having every parameter for these ops duplicated, once as an attribute and once as an SSA value, along with the added custom verifier and assembly format for each op.
Could we consider removing the attributes altogether and instead use SSA values, with arith.constant for the static case? All existing lowerings can get the value from arith.constant and throw an error if it's not a constant, this emitC pass can use the actual SSA values. This approach would of course touch a lot of code (all examples etc. that use these ops with attributes would have to be rewritten to use arith.constant), but I think AI could handle it. I think it would be cleaner and might remove the need for customAssemblyFormat and hasVerifier for every op (haven't gotten to those yet but assume they're there because of this).
| Optionally, SSA values can be provided for 'dyn_address', 'dyn_value', and | ||
| 'dyn_mask' to enable runtime-parameterized sequences. | ||
|
|
||
| Static syntax (unchanged): `aiex.npu.maskwrite32 {address = 123 : ui32, ...}` |
There was a problem hiding this comment.
Suggest removing "(unchanged)" from these comments
| //===- TxnEncoding.h - Standalone TXN instruction encoding -------*- C++ | ||
| //-*-===// |
There was a problem hiding this comment.
suggest formatting onto a single line
| // Use encoding library for the core format, then fix up col/row field. | ||
| aie_runtime::txn_append_blockwrite(instructions, *address, payload.data(), | ||
| payload.size()); | ||
|
|
||
| // XAIE_IO_BLOCKWRITE | ||
| words[0] = XAIE_IO_BLOCKWRITE; | ||
| words[2] = op.getAddress(); | ||
| // The encoding library leaves word[1] as 0. If col/row are present, set it. | ||
| auto col = op.getColumn(); | ||
| auto row = op.getRow(); | ||
| if (col && row) { | ||
| words[1] = (*col & 0xff) | ((*row & 0xff) << 8); | ||
| // word[1] is at position (current_size - headerSize - count + 1) | ||
| size_t headerPos = instructions.size() - 4 - payload.size(); | ||
| instructions[headerPos + 1] = (*col & 0xff) | ((*row & 0xff) << 8); | ||
| } |
There was a problem hiding this comment.
Why doesn't the encoding library aie_runtime::txn_append_blockwrite take row/col?
There was a problem hiding this comment.
I think this file could benefit from some code deduplication and cleanup. If NpuWriteBd op were changed to also accept SSA values, maybe there would be less of a need for separate code paths here. I'd like to avoid having to make every change in two places (dynamic and static path) for future changes to these ops.
| GreedyRewriteConfig rewriter_config = GreedyRewriteConfig(); | ||
| rewriter_config.setRegionSimplificationLevel( | ||
| GreedySimplifyRegionLevel::Disabled); | ||
| rewriter_config.enableConstantCSE(false); |
| //===- AIETargetCppTxn.cpp - EmitC-based C++ TXN translation ------*- C++ | ||
| //-*-===// |
| /// Extract design-specific constants from the static instruction stream. | ||
| /// | ||
| /// The static instructions always begin with: | ||
| /// [4-word TXN header] | ||
| /// [6-word write32: RTP write 0 — rtp_addr = words[header+2]] | ||
| /// [6-word write32: RTP write 1] | ||
| /// ... then DMA configuration including a maskwrite before S2MM push ... | ||
| /// | ||
| /// We scan for the first write32 (opcode 0) to get the RTP address, and | ||
| /// for the first maskwrite32 (opcode 3) to get the S2MM control register | ||
| /// address and its value/mask. | ||
| inline DesignConstants extract_constants(const std::vector<uint32_t> &insts) { | ||
| DesignConstants c{}; | ||
| constexpr uint32_t HEADER_SIZE = 4; | ||
|
|
||
| bool found_rtp = false, found_s2mm = false; | ||
| size_t i = HEADER_SIZE; | ||
| while (i < insts.size() && (!found_rtp || !found_s2mm)) { | ||
| uint32_t opcode = insts[i]; | ||
|
|
||
| if (opcode == aie_runtime::TXN_OPC_WRITE && !found_rtp) { | ||
| // First write32: RTP address is at word [i+2] | ||
| c.rtp_addr = insts[i + 2]; | ||
| found_rtp = true; | ||
| i += 6; | ||
| } else if (opcode == aie_runtime::TXN_OPC_MASKWRITE && !found_s2mm) { | ||
| // First maskwrite: S2MM control register | ||
| c.s2mm_ctrl = insts[i + 2]; | ||
| c.s2mm_ctrl_val = insts[i + 4]; | ||
| c.s2mm_ctrl_mask = insts[i + 5]; | ||
| found_s2mm = true; | ||
| i += 7; | ||
| } else { | ||
| // Skip this op by reading its size field | ||
| uint32_t op_size_bytes = 0; | ||
| if (opcode == aie_runtime::TXN_OPC_WRITE) | ||
| op_size_bytes = insts[i + 5]; | ||
| else if (opcode == aie_runtime::TXN_OPC_MASKWRITE) | ||
| op_size_bytes = insts[i + 6]; | ||
| else if (opcode == aie_runtime::TXN_OPC_BLOCKWRITE) | ||
| op_size_bytes = insts[i + 3]; | ||
| else if (opcode == aie_runtime::TXN_OPC_TCT) | ||
| op_size_bytes = insts[i + 1]; | ||
| else if (opcode == aie_runtime::TXN_OPC_DDR_PATCH) | ||
| op_size_bytes = insts[i + 1]; | ||
| else | ||
| break; // unknown opcode | ||
|
|
||
| i += op_size_bytes / sizeof(uint32_t); | ||
| } | ||
| } | ||
|
|
||
| if (!found_rtp) | ||
| throw std::runtime_error("Could not find RTP write in static instructions"); | ||
| if (!found_s2mm) | ||
| throw std::runtime_error( | ||
| "Could not find S2MM maskwrite in static instructions"); | ||
|
|
||
| return c; | ||
| } |
There was a problem hiding this comment.
Might it make sense to add a compiler option to export some of these "magic values" at compile time, into say a JSON that could be ingested at runtime? Some of them, like the controller ID, probably should also just get baked into the code generated by emitC.
There was a problem hiding this comment.
This doesn't seem to use the code generated from the emitC, but instead reconstructs the instruction sequence manually using the encoding library.
Summary
TxnEncoding.h) for building NPU instruction binaries at runtime--aie-generate-txn-cppcompiler flag: generates C++ code from parameterized MLIR runtime sequences via EmitCnpu.dma_memcpy_ndlowering with SSA sizes/strides/offsetsIsolatedFromAboveonRuntimeSequenceOpfor correct SCF→CF scoping