Skip to content

Fix ktdp_lowering.construct_indirect_access_tile to align with ktdp variant semantics - #140

Merged
bmahjour merged 5 commits into
torch-spyre:mainfrom
bmahjour:fix-ind-access-tile
Sep 10, 2026
Merged

Fix ktdp_lowering.construct_indirect_access_tile to align with ktdp variant semantics#140
bmahjour merged 5 commits into
torch-spyre:mainfrom
bmahjour:fix-ind-access-tile

Conversation

@bmahjour

@bmahjour bmahjour commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Fix ktdp_lowering.construct_indirect_access_tile to align with ktdp variant semantics

What was wrong

The op was structurally inconsistent with its ktdp.construct_indirect_access_tile counterpart in several ways:

Semantic / structural issues:

  • intermediate_variables were modelled as plain Variadic<Index> operands read from the enclosing scope. In the ktdp variant they are block arguments of a hidden single-block region — values defined by the op, not passed into it. This distinction matters: they enumerate the iteration space of the tile and must not be SSA-reachable as ordinary operands.
  • Direct subscripts into $base were Variadic<Index> (flat SSA values). This cannot represent affine expressions like (%c0 + %arg7) which appear throughout the transformation pipeline without pre-materializing intermediate arith ops. The ktdp variant stores these as ArrayAttr of AffineMapAttr over a unified (captured..., intermediate...) domain for exactly this reason.
  • IAB subscripts were also Variadic<Index> operands. While the representation is correct in principle (IAB subscripts are always bare variable references because spyreop.idx32toaddr pre-computes the flat address), listing them as operands was wrong because intermediate-variable IAB subscripts are region block args, not outer SSA values — resolving them as operands would fail at parse time.

Implementation issues (discovered during parser/printer audit against the ktdp variant):

  • Parser built the hidden region before canonicalization; ktdp builds it after.
  • ensureTerminator was called after adding block arguments; ktdp calls it before.
  • Map canonicalization used expr.replaceDims per result; ktdp uses raw.compose(remap).
  • No infinite-loop guard in the subscript parse loop.
  • Builder lacked the three asserts that guard variables_space_set / variables_space_order dimension consistency.
  • Printer used llvm::interleaveComma for block args; ktdp uses p << getRegion().getArguments().

What was changed

Op structure (KTDPLowering.td):

  • Removed Variadic<Index>:$intermediate_variables from arguments and replaced with let regions = (region SizedRegion<1>:$region) + SingleBlockImplicitTerminator<"ktdp::RegionTerminatorOp">, exactly mirroring the ktdp variant.
  • Replaced Variadic<Index>:$direct_subscripts with ArrayAttr:$per_dim_subscript_maps — one AffineMapAttr per $base dimension over the unified (captured..., intermediate...) domain.
  • Replaced Variadic<Index>:$ind_addr_buf_subscripts with DenseI32ArrayAttr:$ind_addr_buf_dim_positions — one integer per IAB dimension giving its position in the unified variable list. No affine map is needed here because IAB entries already hold flat absolute addresses; the subscript is structurally always a bare variable reference.
  • Removed AttrSizedOperandSegments (now only one variadic operand segment: $captured_variables).
  • Added getIntermediateVariables() helper via extraClassDeclaration, mirroring the ktdp variant.

Implementation (KTDPLoweringOps.cpp):

  • Parser: region built after all parsing/canonicalization with result.regions.reserve(1); ensureTerminator called before block arguments; map canonicalization switched to raw.compose(remap); posMap replaces redundant dual-scan lookup; infinite-loop guard added.
  • Builder: ensureTerminator called before block arguments; three asserts added for variables_space_set / variables_space_order / numIntermediateVariables consistency.
  • Printer: p << getRegion().getArguments() for block-arg output.

Tests (indirect_access_tile.mlir):

  • Rewrote all test cases: function signatures no longer carry intermediate variables as arguments; subscripts use (affine-expr) syntax; a fourth test case added covering the post-per-entry-legalization form where the IAB subscript is a captured loop IV.

@KFAFSP

KFAFSP commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

I think this warrants a verifier that checks the block is empty (excluding the terminator), at the very least.

Honestly, I find the whole "hidden block just to define SSA values" super weird to begin with. We could just add the intermediate variables to the domain of the affine map and let the user specify the map directly. With a bit of finesse, we could even try and keep the syntax - as long as you don't resolveOperand, I think you can just eat the SSA ids.

…ion and behave more consistently with its ktdp variant

Signed-off-by: Bardia Mahjour <bmahjour@ca.ibm.com>
Signed-off-by: Bardia Mahjour <bmahjour@ca.ibm.com>
@bmahjour
bmahjour force-pushed the fix-ind-access-tile branch from 4e421d3 to 3ee25b6 Compare September 3, 2026 20:10
@bmahjour

bmahjour commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator Author

I think this warrants a verifier that checks the block is empty (excluding the terminator), at the very least.

Honestly, I find the whole "hidden block just to define SSA values" super weird to begin with. We could just add the intermediate variables to the domain of the affine map and let the user specify the map directly. With a bit of finesse, we could even try and keep the syntax - as long as you don't resolveOperand, I think you can just eat the SSA ids.

Again I see your point, but since the upstream variant is already in this form, I strongly suggest we keep consistent until such time that we revisit the KTIR syntax.

Comment thread include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td Outdated
Comment thread include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td Outdated
…_buf_memref

Signed-off-by: Bardia Mahjour <bmahjour@ca.ibm.com>
@bmahjour
bmahjour requested review from KFAFSP and nkhoun September 4, 2026 14:24
Comment thread include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Signed-off-by: Bardia Mahjour <bmahjour@ca.ibm.com>
@bmahjour
bmahjour requested a review from KFAFSP September 9, 2026 18:30
Comment thread lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp Outdated
Signed-off-by: Bardia Mahjour <bmahjour@ca.ibm.com>
@bmahjour
bmahjour requested a review from KFAFSP September 9, 2026 21:55
@bmahjour
bmahjour merged commit 4bb4c51 into torch-spyre:main Sep 10, 2026
4 checks passed
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