Skip to content

Runtime TXN generation library and compiler support#3002

Draft
jgmelber wants to merge 19 commits intomainfrom
dynamic-runtime-sequences
Draft

Runtime TXN generation library and compiler support#3002
jgmelber wants to merge 19 commits intomainfrom
dynamic-runtime-sequences

Conversation

@jgmelber
Copy link
Copy Markdown
Collaborator

@jgmelber jgmelber commented Mar 26, 2026

Summary

  • Standalone TXN encoding library (TxnEncoding.h) for building NPU instruction binaries at runtime
  • --aie-generate-txn-cpp compiler flag: generates C++ code from parameterized MLIR runtime sequences via EmitC
  • Dynamic npu.dma_memcpy_nd lowering with SSA sizes/strides/offsets
  • IsolatedFromAbove on RuntimeSequenceOp for correct SCF→CF scoping
  • Single-core dynamic GEMM example: one XCLBIN, any M/K/N at runtime

jgmelber and others added 15 commits March 26, 2026 08:11
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]>
@jgmelber jgmelber force-pushed the dynamic-runtime-sequences branch from 2c839d4 to 93acc43 Compare March 26, 2026 19:36
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]>
@jgmelber jgmelber force-pushed the dynamic-runtime-sequences branch from 47bb7b4 to c97ce07 Compare March 26, 2026 22:06
@jgmelber jgmelber force-pushed the dynamic-runtime-sequences branch from 3e32ea1 to 5286ec4 Compare March 26, 2026 22:24
@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Mar 26, 2026

Coverage Report

Created: 2026-03-26 23:50

Click here for information about interpreting this report.

FilenameFunction CoverageLine CoverageRegion CoverageBranch Coverage
Conversion/AIEToConfiguration/AIEToConfiguration.cpp 91.30% 67.58% 61.27% 45.19%
Conversion/AIEXToEmitC/AIEXToEmitC.cpp 5.00% 0.38% 0.35% 0.00%
Dialect/AIE/Transforms/AIEInsertTraceFlows.cpp 100.00% 90.78% 81.71% 78.48%
Dialect/AIE/Transforms/AIEVectorTransferLowering.cpp 83.33% 79.17% 72.73% 50.00%
Dialect/AIEX/IR/AIEXDialect.cpp 98.36% 76.45% 71.49% 54.81%
Dialect/AIEX/Transforms/AIEDMATasksToNPU.cpp 95.45% 85.14% 89.30% 79.65%
Dialect/AIEX/Transforms/AIEDmaToNpu.cpp 100.00% 82.62% 76.28% 54.76%
Dialect/AIEX/Transforms/AIELowerSetLock.cpp 100.00% 82.35% 80.00% 50.00%
Dialect/AIEX/Transforms/AIEMaterializeBDChains.cpp 100.00% 84.71% 80.00% 57.14%
Totals 86.74% 65.32% 63.41% 49.75%
Generated by llvm-cov -- llvm version 18.1.3

@jgmelber jgmelber force-pushed the dynamic-runtime-sequences branch from 53af3d9 to 934d90b Compare March 26, 2026 22:29
Comment on lines +275 to +276
f.write(
"""
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
f.write(
"""
f.write("""

Comment on lines +283 to +284
"""
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
"""
)
""")

Comment on lines +110 to +111
f.write(
"""extern "C" {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
f.write(
"""extern "C" {
f.write("""extern "C" {

Comment on lines +117 to +118
}"""
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
}"""
)
}""")

Comment on lines +188 to +189
f.write(
"""extern "C" {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
f.write(
"""extern "C" {
f.write("""extern "C" {

Comment on lines +193 to +194
}"""
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[black] reported by reviewdog 🐶

Suggested change
}"""
)
}""")

Copy link
Copy Markdown
Collaborator

@andrej andrej left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2026

Comment on lines +751 to 764
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; }
}];
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, ...}`
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggest removing "(unchanged)" from these comments

Comment on lines +1 to +2
//===- TxnEncoding.h - Standalone TXN instruction encoding -------*- C++
//-*-===//
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggest formatting onto a single line

Comment on lines +109 to 120
// 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);
}
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why doesn't the encoding library aie_runtime::txn_append_blockwrite take row/col?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this needed?

Comment on lines +1 to +2
//===- AIETargetCppTxn.cpp - EmitC-based C++ TXN translation ------*- C++
//-*-===//
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

format

Comment on lines +68 to +127
/// 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;
}
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doesn't seem to use the code generated from the emitC, but instead reconstructs the instruction sequence manually using the encoding library.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants