diff --git a/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.h b/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.h index fecc3d86..0523c66b 100644 --- a/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.h +++ b/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.h @@ -31,7 +31,10 @@ #include #include "dataflow-scheduler/Dialect/KTDPLowering/KTDPLoweringDialect.h" // IWYU pragma: keep -#include "ktir/Dialect/KTDP/KTDPTypes.h" +// Full KTDP ops header needed so that ktdp::RegionTerminatorOp (used as the +// SingleBlockImplicitTerminator of our hidden region) is declared before the +// generated .h.inc is parsed by the C++ compiler. +#include "ktir/Dialect/KTDP/KTDP.h" /// Auto-generated includes. #define GET_OP_CLASSES diff --git a/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td b/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td index 08fba647..cf925fc8 100644 --- a/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td +++ b/include/dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.td @@ -32,6 +32,9 @@ include "mlir/Interfaces/ViewLikeInterface.td" include "dataflow-scheduler/Dialect/KTDPLowering/KTDPLoweringDialect.td" // Import the KTDP types so we can use Ktdp_AccessTileType as the result type. +// ktdp::RegionTerminatorOp is referenced by its C++ class name only (as a +// string argument to SingleBlockImplicitTerminator), so we do not need to +// include the full KTDP ops tablegen here. include "ktir/Dialect/KTDP/KTDPTypes.td" class KTDPLowering_Op traits = []> : @@ -43,7 +46,7 @@ class KTDPLowering_Op traits = []> : def KTDPLowering_ConstructIndirectAccessTileOp : KTDPLowering_Op<"construct_indirect_access_tile", - [AttrSizedOperandSegments]> { + [SingleBlockImplicitTerminator<"ktdp::RegionTerminatorOp">]> { let summary = "indirect access tile with indirection factored into an IAB " "base address and direct element offsets"; @@ -51,88 +54,147 @@ def KTDPLowering_ConstructIndirectAccessTileOp : The `ktdp_lowering.construct_indirect_access_tile` operation is the lowered form of `ktdp.construct_indirect_access_tile` in which the indirect address has been factored: the base address contribution is - supplied by an entry in the indirect address buffer (IAB), selected by - `$ind_addr_buf_subscripts` into `$ind_addr_buf_memref`, and the remaining - dimensions are expressed as direct element offsets into `$base` via - `$direct_subscripts`. In all other respects the semantics of this op are - inherited from `ktdp.construct_indirect_access_tile`. - - **Semantics of `base_ptr = %iab_mv[%s0, %s1, ...]`:** the `index` value - stored at those subscripts in the IAB memref is added as an element-offset - delta to the flat address computed from `$base`'s `construct_memory_view` - base together with the explicit `$direct_subscripts`. Concretely, if the - `ind()` dimension of `$base` had stride `S`, the IAB entry stores - `original_index * S`, and the explicit subscript for that dimension is - `%c0` (contributing `0 * S = 0`). The two terms together reproduce the - original indirect address without double-counting. - - Any of the `$ind_addr_buf_subscripts` may be values drawn from the - `$intermediate_variables` list, which are the same block arguments - that parameterise `$direct_subscripts` and `$variables_space_set`. - - The `$intermediate_variables` and `$variables_space_set` define the tile - shape. `ConstructThreeStagePipeline` reads `$variables_space_set` to - determine the element count of the `ktdf.ind_data_transfer` it emits. + supplied by a pre-computed flat address stored in an entry of the indirect + address buffer (IAB), selected by `$ind_addr_buf_subscripts` into + `$ind_addr_buf_memref`. The remaining dimensions of `$base` are described + by `$per_dim_subscript_maps` — one affine map per dimension — in all other + respects the semantics of this op follow those of + `ktdp.construct_indirect_access_tile`. + + **Intermediate variables** are the loop-iteration values that enumerate + points in the variable space defined by `$variables_space_set`. They are + modelled as block arguments of the op's hidden single-block region — exactly + as in the `ktdp` variant — and are *not* passed as operands from the + enclosing scope. Access them via `getIntermediateVariables()`. + + **IAB subscripts vs. direct subscript maps:** + + `$ind_addr_buf_subscripts` are plain SSA `index` values — each one is + either a block argument of `$region` (an intermediate variable) or a + member of `$captured_variables` (an outer loop IV or constant). No affine + map is needed here because the subscript is used directly as an index into + `$ind_addr_buf_memref`; the address buffer lookup itself supplies the + selected address, so the subscript only needs to identify which entry to + read. Note that even at the KTDP level, an affine map is not applied to + the value loaded from an indirect subscript either. + + `$per_dim_subscript_maps` are affine maps over the unified dimension space + `(captured_variables..., intermediate_variables...)`. They *can* contain + arithmetic — e.g. `(%c0 + %arg7)` lowers to `affine_map<(d0,d1)->d0+d1>` + applied to `(%c0, %arg7)` — because the hardware memory unit still applies + strides to these subscripts at transfer time. The canonical dimension + ordering mirrors `ktdp.construct_indirect_access_tile`: + captured variables occupy the leading dimensions, intermediate variables + follow. + + **Semantics of `base_ptr = %iab_mv[%s0, %s1, ...]`:** the flat address + value stored at those subscripts in the IAB memref is added as an + element-offset delta to the address computed from `$base`'s + `construct_memory_view` base. The IAB value already represents the + complete contribution of the indirect dimensions; no stride or affine + subscript map is applied to an indirect subscript or to the value loaded + from it. `$per_dim_subscript_maps` contribute only the coordinates for + direct dimensions, whose strides are applied by the hardware memory unit: + `flat = cmv_base + iab[s0,s1,...] + direct_map0(...)*S0 + ...` + + The `$variables_space_set` defines the tile shape. + `ConstructThreeStagePipeline` reads it to determine the element count of + the `ktdf.ind_data_transfer` it emits. The result type `!ktdp.access_tile<...>` is identical to the type produced by `ktdp.construct_indirect_access_tile`, so all existing consumers (`ktdp.load`, `ktdp.store`, `ConstructThreeStagePipeline`) work unchanged. - Example (after ind_addr_buf capacity legalization, 1-D IAB; - `%arg6` is an intermediate variable used as the IAB subscript): + Example (after per-entry legalization, 1-D IAB; `%i2` is a captured + loop IV selecting the IAB entry; `%arg7` and `%arg8` are intermediate + variables; `%c0` is a captured constant): ```mlir + // per_dim_subscript_maps: + // dim 0 (indirect dim, zero'd out): affine_map<(d0,d1,d2) -> (0)> + // dim 1 (direct, offset by captured c0): affine_map<(d0,d1,d2) -> (d0 + d1)> + // dim 2 (direct, identity): affine_map<(d0,d1,d2) -> (d2)> + // captured_variables: [%c0, %arg7] intermediate_variables: [%arg8] %tile = ktdp_lowering.construct_indirect_access_tile - intermediate_variables(%arg6, %arg7, %arg8) - base_ptr = %iab_mv[%arg6] - %X[%c0, %arg7, %arg8] + intermediate_variables(%arg7, %arg8) + base_ptr = %iab_mv[%i2] + %desc_1[(%c0), (%c0 + %arg7), (%arg8)] {variables_space_set = #set, variables_space_order = #map} - : memref<64x2x64xf16>, memref<32xindex, "IAB"> -> !ktdp.access_tile<32x2x64xindex> + : memref<64x2x64xf16>, memref<32xindex, "IAB"> + -> !ktdp.access_tile<2x64xindex> ``` - Example (before capacity legalization, 2-D IAB; - `%arg5` and `%arg6` are intermediate variables): + Example (before capacity legalization, 2-D IAB; all four intermediate + variables are block args): ```mlir %tile = ktdp_lowering.construct_indirect_access_tile intermediate_variables(%arg5, %arg6, %arg7, %arg8) base_ptr = %iab_mv[%arg5, %arg6] - %X[%c0, %arg7, %arg8] + %desc_1[(%c0), (%c0 + %arg7), (%arg8)] {variables_space_set = #set2, variables_space_order = #map} - : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -> !ktdp.access_tile<2x32x2x64xindex> + : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> + -> !ktdp.access_tile<2x32x2x64xindex> ``` }]; let arguments = (ins AnyMemRef:$base, // IAB memory view (the memory space is architecture-defined; - // the verifier only checks that ind_addr_buf_subscripts count - // equals the rank of this memref). + // the verifier checks that ind_addr_buf_dim_positions.size() == rank of this + // memref). Arg:$ind_addr_buf_memref, - // Per-dimension subscripts into ind_addr_buf_memref; count must equal - // the rank of ind_addr_buf_memref (enforced by the verifier). - // Any of these may be values drawn from $intermediate_variables. - Variadic:$ind_addr_buf_subscripts, - // Explicit direct subscripts into $base (one per base dimension). - Variadic:$direct_subscripts, - // Intermediate variables that parameterise the tile shape via - // variables_space_set; may also appear in ind_addr_buf_subscripts - // and direct_subscripts. These operands play the same role as - // the hidden-region block arguments used by the ktdp variant: they - // represent the loop-like iteration variables that enumerate points - // in the variable space defined by variables_space_set. - Variadic:$intermediate_variables, + // One integer per IAB dimension giving its position in the unified + // (captured_variables..., intermediate_variables...) ordering. No operand + // is needed because each IAB subscript is always a direct use of exactly + // one variable from that space — never an arithmetic expression — so a + // simple index into the unified list is sufficient. + DenseI32ArrayAttr:$ind_addr_buf_dim_positions, + // One AffineMapAttr per dimension of $base. Each map is expressed over + // the unified dimension space (captured_variables..., intermediate_variables...). + // Maps CAN contain arithmetic (e.g. d0 + d1, d2 mod 64) because the + // hardware memory unit still applies strides to these values at transfer time. + ArrayAttr:$per_dim_subscript_maps, + // SSA values from the enclosing scope referenced by $per_dim_subscript_maps + // that are NOT intermediate variables. + // These occupy the leading dimensions of every subscript map. + Variadic:$captured_variables, AffineMapAttr:$variables_space_order, Builtin_IntegerSetAttr:$variables_space_set ); + // Hidden single-block region whose block arguments are the intermediate + // variables (the loop-iteration indices that enumerate the variable space). + // The block is terminated by ktdp.region_terminator. + let regions = (region SizedRegion<1>:$region); + let results = (outs Ktdp_AccessTileType:$result); let hasVerifier = 1; - let assemblyFormat = [{ - `intermediate_variables` `(` $intermediate_variables `)` - `base_ptr` `=` $ind_addr_buf_memref `[` $ind_addr_buf_subscripts `]` - $base `[` $direct_subscripts `]` - attr-dict `:` type($base) `,` type($ind_addr_buf_memref) `->` qualified(type($result)) + let skipDefaultBuilders = 1; + let hasCustomAssemblyFormat = 1; + + let builders = [ + OpBuilder<(ins + "::mlir::ktdp::AccessTileType":$resultType, + "::mlir::Value":$base, + "::mlir::Value":$indAddrBufMemref, + "::mlir::DenseI32ArrayAttr":$indAddrBufDimPositions, + "::mlir::ArrayAttr":$perDimSubscriptMaps, + "::mlir::ValueRange":$capturedVariables, + "unsigned":$numIntermediateVariables, + "::mlir::AffineMap":$variablesSpaceOrder, + "::mlir::IntegerSet":$variablesSpaceSet)> + ]; + + let extraClassDeclaration = [{ + /// Returns the intermediate variables (block arguments of the hidden + /// region), which enumerate the variable space defined by + /// $variables_space_set. Mirrors + /// ConstructIndirectAccessTilesOp::getIntermediateVariables() in the + /// ktdp dialect. + ::mlir::ValueRange getIntermediateVariables() { + return getRegion().getArguments(); + } }]; } diff --git a/lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp b/lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp index 33b57fd0..255f3121 100644 --- a/lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp +++ b/lib/Dialect/KTDPLowering/KTDPLoweringOps.cpp @@ -24,6 +24,7 @@ #include "dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.h" // clang-format on +#include #include #include #include @@ -50,18 +51,315 @@ void KTDPLoweringDialect::registerOps() { #define GET_OP_CLASSES #include "dataflow-scheduler/Dialect/KTDPLowering/KTDPLowering.cpp.inc" +//===----------------------------------------------------------------------===// +// ConstructIndirectAccessTileOp — builder +//===----------------------------------------------------------------------===// + +// Build method contract +// --------------------- +// `perDimSubscriptMaps` must already be canonicalized over the unified +// dimension ordering (capturedVariables..., intermediateVariables...). +// `indAddrBufDimPositions[i]` is the index of the i-th IAB dimension in that +// same unified ordering. +// `numIntermediateVariables` must equal variablesSpaceSet.getNumDims(). +void ConstructIndirectAccessTileOp::build( + OpBuilder& builder, OperationState& result, ktdp::AccessTileType resultType, + Value base, Value indAddrBufMemref, + DenseI32ArrayAttr indAddrBufDimPositions, ArrayAttr perDimSubscriptMaps, + ValueRange capturedVariables, unsigned numIntermediateVariables, + AffineMap variablesSpaceOrder, IntegerSet variablesSpaceSet) { + assert( + variablesSpaceSet.getNumDims() == variablesSpaceOrder.getNumInputs() && + "variables_space_order input count must match variables_space_set dims"); + assert(variablesSpaceOrder.getNumInputs() == + variablesSpaceOrder.getNumResults() && + "variables_space_order must have equal input and output dimensions"); + assert(variablesSpaceSet.getNumDims() == numIntermediateVariables && + "numIntermediateVariables must equal variables_space_set dims"); + + result.addOperands(base); + result.addOperands(indAddrBufMemref); + result.addOperands(capturedVariables); + auto& props = result.getOrAddProperties(); + props.ind_addr_buf_dim_positions = indAddrBufDimPositions; + props.per_dim_subscript_maps = perDimSubscriptMaps; + props.variables_space_order = AffineMapAttr::get(variablesSpaceOrder); + props.variables_space_set = IntegerSetAttr::get(variablesSpaceSet); + + // Hidden region: one index-typed block arg per intermediate variable. + // ensureTerminator is called before adding block arguments, matching the + // pattern established by ConstructIndirectAccessTilesOp in the ktdp dialect. + Region* region = result.addRegion(); + Block& body = region->emplaceBlock(); + ensureTerminator(*region, builder, builder.getUnknownLoc()); + for (unsigned i = 0; i < numIntermediateVariables; ++i) + body.addArgument(builder.getIndexType(), builder.getUnknownLoc()); + + result.types.push_back(resultType); +} + +//===----------------------------------------------------------------------===// +// ConstructIndirectAccessTileOp — custom assembly format +// +// Printed form: +// +// ktdp_lowering.construct_indirect_access_tile +// intermediate_variables(%iv0, %iv1, ...) +// base_ptr = %iab[%var_at_pos0, %var_at_pos1, ...] +// %base[(affine-expr), (affine-expr), ...] +// { attr-dict } +// : type($base), type($ind_addr_buf_memref) -> qualified(type($result)) +// +// `intermediate_variables` names the hidden region's block arguments. +// IAB subscripts are printed as the SSA names from the unified variable list +// (captured..., intermediate...) at the positions stored in +// `ind_addr_buf_dim_positions`. +// Per-dimension base subscripts use `(affine-expr)` syntax, with maps over the +// unified dimension space. +//===----------------------------------------------------------------------===// + +ParseResult ConstructIndirectAccessTileOp::parse(OpAsmParser& parser, + OperationState& result) { + auto& builder = parser.getBuilder(); + MLIRContext* ctx = builder.getContext(); + + // --- intermediate_variables(%iv0, %iv1, ...) --- + // These will become block arguments of the hidden region; NOT op operands. + SmallVector ivNames; + if (parser.parseKeyword("intermediate_variables") || + parser.parseOperandList(ivNames, AsmParser::Delimiter::Paren)) + return failure(); + + // --- base_ptr = %iab_memref[%var0, %var1, ...] --- + // The IAB subscripts are SSA names drawn from the unified variable space + // (captured + intermediate). We collect their names here and resolve their + // positions after we know the full unified ordering. + OpAsmParser::UnresolvedOperand iabMemref; + SmallVector iabSubscriptNames; + if (parser.parseKeyword("base_ptr") || parser.parseEqual() || + parser.parseOperand(iabMemref) || parser.parseLSquare() || + parser.parseOperandList(iabSubscriptNames) || parser.parseRSquare()) + return failure(); + + // --- %base[affine-expr, ...] --- + OpAsmParser::UnresolvedOperand base; + if (parser.parseOperand(base)) return failure(); + + SmallVector rawMaps; + SmallVector> rawMapOperands; + if (parser.parseCommaSeparatedList(AsmParser::Delimiter::Square, [&]() { + SmallVector dimOps; + SmallVector symOps; + AffineExpr expr; + if (parser.parseAffineExprOfSSAIds(dimOps, symOps, expr)) + return failure(); + + SmallVector mapOps(dimOps); + mapOps.append(symOps.begin(), symOps.end()); + rawMaps.push_back(AffineMapAttr::get( + AffineMap::get(dimOps.size(), symOps.size(), expr, ctx))); + rawMapOperands.push_back(mapOps); + return success(); + })) + return failure(); + result.attributes.clear(); // remove temporaries from parseAffineMapOfSSAIds + + // --- optional attr-dict --- + if (parser.parseOptionalAttrDict(result.attributes)) return failure(); + + // --- : type($base), type($iab) -> type($result) --- + Type baseType, iabType, resultType; + if (parser.parseColon() || parser.parseType(baseType) || + parser.parseComma() || parser.parseType(iabType) || parser.parseArrow() || + parser.parseType(resultType)) + return failure(); + + // --- Determine the unified variable ordering (captured..., intermediate...) + // --- Captured variables: SSA names referenced in per-dim maps or IAB + // subscripts that are NOT in the intermediate-variables list. + llvm::SmallSetVector ivNameSV; + for (auto& iv : ivNames) ivNameSV.insert(iv.name); + + SmallVector capturedNames; + llvm::SmallDenseSet capturedSeen; + auto maybeCapture = [&](StringRef name) { + if (!ivNameSV.contains(name) && capturedSeen.insert(name).second) + capturedNames.push_back(name); + }; + for (auto& ops : rawMapOperands) + for (auto& op : ops) maybeCapture(op.name); + for (auto& sub : iabSubscriptNames) maybeCapture(sub.name); + + unsigned unifiedDims = capturedNames.size() + ivNameSV.size(); + + // Build a position map from name → unified dim index. + llvm::SmallDenseMap posMap; + for (unsigned c = 0; c < capturedNames.size(); ++c) + posMap[capturedNames[c]] = c; + for (unsigned v = 0; v < ivNameSV.size(); ++v) + posMap[ivNameSV[v]] = capturedNames.size() + v; + + // --- Canonicalize per-dim subscript maps to the unified dimension space --- + // Strategy: build a "remap" map from the unified domain to each map's local + // domain, then compose — matching the approach in + // canonicalizeAffineMapsToUnifiedOperands used by the ktdp variant. + SmallVector canonicalMaps; + for (size_t i = 0; i < rawMaps.size(); ++i) { + AffineMap raw = rawMaps[i].getValue(); + // Build remap: (unified dims...) -> (local dim for each map operand) + SmallVector remapResults; + for (auto& opName : rawMapOperands[i]) { + auto it = posMap.find(opName.name); + if (it == posMap.end()) + return parser.emitError(parser.getNameLoc()) + << "subscript operand '" << opName.name + << "' is neither a captured variable nor an intermediate " + "variable"; + remapResults.push_back(getAffineDimExpr(it->second, ctx)); + } + AffineMap remap = + AffineMap::get(unifiedDims, /*numSymbols=*/0, remapResults, ctx); + // raw(localDims...) . remap(unifiedDims...) => canonical(unifiedDims...) + canonicalMaps.push_back(AffineMapAttr::get(raw.compose(remap))); + } + auto& props = result.getOrAddProperties(); + props.per_dim_subscript_maps = builder.getArrayAttr(canonicalMaps); + + // Build the hidden region with one block arg per intermediate variable. + // Region is added after parsing is complete, matching the ktdp variant's + // ordering. ensureTerminator is called before adding block arguments. + Region* region = result.addRegion(); + Block& body = region->emplaceBlock(); + ensureTerminator(*region, builder, result.location); + for (size_t i = 0; i < ivNames.size(); ++i) + body.addArgument(builder.getIndexType(), builder.getUnknownLoc()); + + // --- Resolve IAB subscript names to positions in the unified ordering --- + SmallVector iabPositions; + for (auto& sub : iabSubscriptNames) { + auto it = posMap.find(sub.name); + if (it == posMap.end()) + return parser.emitError(parser.getNameLoc()) + << "IAB subscript '" << sub.name + << "' is neither a captured variable nor an intermediate variable"; + iabPositions.push_back(static_cast(it->second)); + } + props.ind_addr_buf_dim_positions = builder.getDenseI32ArrayAttr(iabPositions); + + // --- Resolve operands --- + if (parser.resolveOperand(base, baseType, result.operands) || + parser.resolveOperand(iabMemref, iabType, result.operands)) + return failure(); + + // Collect and resolve captured variables in first-occurrence order. + SmallVector capturedResolvable; + { + llvm::SmallDenseSet resolvedSeen; + auto collect = [&](llvm::ArrayRef ops) { + for (auto& op : ops) + if (!ivNameSV.count(op.name) && resolvedSeen.insert(op.name).second) + capturedResolvable.push_back(op); + }; + for (auto& ops : rawMapOperands) collect(ops); + collect(iabSubscriptNames); + } + if (parser.resolveOperands(capturedResolvable, builder.getIndexType(), + result.operands)) + return failure(); + + result.addTypes(resultType); + return success(); +} + +void ConstructIndirectAccessTileOp::print(OpAsmPrinter& p) { + // Build the unified value list once: (captured..., intermediate...) + SmallVector allVars(getCapturedVariables().begin(), + getCapturedVariables().end()); + for (Value iv : getIntermediateVariables()) allVars.push_back(iv); + + // intermediate_variables(...) — region block args. + // Use p << args to match the ktdp variant's idiomatic block-arg printing. + p << " intermediate_variables("; + p << getRegion().getArguments(); + p << ")"; + + // base_ptr = %iab[%var_at_pos0, ...] + p << " base_ptr = " << getIndAddrBufMemref() << "["; + llvm::interleaveComma(getIndAddrBufDimPositions(), p, [&](int32_t pos) { + p << allVars[static_cast(pos)]; + }); + p << "]"; + + // %base[affine-expr, ...] + p << " " << getBase() << "["; + auto maps = getPerDimSubscriptMaps(); + for (unsigned i = 0, e = maps.size(); i < e; ++i) { + if (i > 0) p << ", "; + p.printAffineMapOfSSAIds(llvm::cast(maps[i]), allVars); + } + p << "]"; + + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{getIndAddrBufDimPositionsAttrName(), + getPerDimSubscriptMapsAttrName()}); + + p << " : " << getBase().getType() << ", " << getIndAddrBufMemref().getType() + << " -> "; + p.printType(getResult().getType()); +} + //===----------------------------------------------------------------------===// // ConstructIndirectAccessTileOp — verifier //===----------------------------------------------------------------------===// LogicalResult ConstructIndirectAccessTileOp::verify() { - // ind_addr_buf_subscripts count must equal the rank of ind_addr_buf_memref. - auto iab_type = mlir::cast(getIndAddrBufMemref().getType()); - auto num_subscripts = static_cast(getIndAddrBufSubscripts().size()); - if (num_subscripts != iab_type.getRank()) - return emitOpError() << "ind_addr_buf_subscripts has " << num_subscripts - << " operand(s) but ind_addr_buf_memref has rank " - << iab_type.getRank() << "; they must be equal"; + // The region must be empty (only the implicit terminator is allowed). + Block& body = getRegion().front(); + if (!body.without_terminator().empty()) + return emitOpError("region must be empty (only the terminator is allowed)"); + + auto iabType = mlir::cast(getIndAddrBufMemref().getType()); + unsigned unifiedDims = + getCapturedVariables().size() + getIntermediateVariables().size(); + + // IAB dim-positions count must match the rank of the IAB memref. + auto iabPositions = getIndAddrBufDimPositions(); + if (static_cast(iabPositions.size()) != iabType.getRank()) + return emitOpError() << "ind_addr_buf_dim_positions has " + << iabPositions.size() + << " entry/entries but ind_addr_buf_memref has rank " + << iabType.getRank() << "; they must be equal"; + + // Each IAB position must be within [0, unifiedDims). + for (auto [idx, pos] : llvm::enumerate(iabPositions)) { + if (pos < 0 || static_cast(pos) >= unifiedDims) + return emitOpError() << "ind_addr_buf_dim_positions[" << idx + << "] = " << pos << " is out of range [0, " + << unifiedDims << ")"; + } + + // per_dim_subscript_maps count must match the rank of $base. + auto baseType = mlir::cast(getBase().getType()); + auto numMaps = static_cast(getPerDimSubscriptMaps().size()); + if (numMaps != baseType.getRank()) + return emitOpError() << "per_dim_subscript_maps has " << numMaps + << " map(s) but base memref has rank " + << baseType.getRank() << "; they must be equal"; + + // Each map must have exactly `unifiedDims` input dimensions. + for (auto [idx, attr] : llvm::enumerate(getPerDimSubscriptMaps())) { + auto mapAttr = mlir::dyn_cast(attr); + if (!mapAttr) + return emitOpError() << "per_dim_subscript_maps[" << idx + << "] is not an AffineMapAttr"; + if (mapAttr.getValue().getNumDims() != unifiedDims) + return emitOpError() + << "per_dim_subscript_maps[" << idx << "] has " + << mapAttr.getValue().getNumDims() + << " dimension(s) but the unified variable space has " + << unifiedDims << " (captured=" << getCapturedVariables().size() + << ", intermediate=" << getIntermediateVariables().size() << ")"; + } return success(); } diff --git a/test/Dialect/KTDPLowering/indirect_access_tile.mlir b/test/Dialect/KTDPLowering/indirect_access_tile.mlir index 10e1b852..8b52f8f3 100644 --- a/test/Dialect/KTDPLowering/indirect_access_tile.mlir +++ b/test/Dialect/KTDPLowering/indirect_access_tile.mlir @@ -3,110 +3,156 @@ // Verifies round-trip parsing and printing of // ktdp_lowering.construct_indirect_access_tile. -// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1) -> (d0, d1)> -// CHECK-DAG: #[[SET1:.*]] = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 63 >= 0)> -// CHECK-DAG: #[[SET2:.*]] = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)> +// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> +// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK-DAG: #[[SET3D:.*]] = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 63 >= 0)> +// CHECK-DAG: #[[SET2D:.*]] = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)> -#set1 = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 63 >= 0)> -#set2 = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)> -#map = affine_map<(d0, d1) -> (d0, d1)> +// Affine maps for the per-dim subscripts used across tests. +// Unified dim ordering: (captured..., intermediate...). +// +// 3-dim base (memref<64x2x64xf16>) with 1 captured (%c0) + 2 iv (%arg7, %arg8): +// dim 0 → %c0 : affine_map<(d0,d1,d2) -> (d0)> +// dim 1 → %c0 + %arg7 : affine_map<(d0,d1,d2) -> (d0 + d1)> +// dim 2 → %arg8 : affine_map<(d0,d1,d2) -> (d2)> +#set3d = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 31 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#set2d = affine_set<(d0, d1) : (d0 >= 0, -d0 + 31 >= 0, d1 >= 0, -d1 + 63 >= 0)> +#map2 = affine_map<(d0, d1) -> (d0, d1)> +#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// ----------------------------------------------------------------------- +// Test 1: 2-D IAB, both IAB subscripts are intermediate variables. +// +// After IndirectComputeGroupSplit (before capacity legalization): +// - 4 intermediate vars: %arg5, %arg6 (IAB dims), %arg7, %arg8 (direct dims) +// - 1 captured: %c0 +// - per_dim_subscript_maps (unified dims = 1 captured + 4 iv = 5): +// dim 0: affine_map<(d0,d1,d2,d3,d4) -> (d0)> (%c0) +// dim 1: affine_map<(d0,d1,d2,d3,d4) -> (d0 + d3)> (%c0 + %arg7) +// dim 2: affine_map<(d0,d1,d2,d3,d4) -> (d4)> (%arg8) +// ----------------------------------------------------------------------- // CHECK-LABEL: func @roundtrip_2d_iab( -// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> -// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<2x32xindex, "IAB"> -// CHECK-SAME: [[C0:%arg[0-9]+]]: index -// CHECK-SAME: [[S0:%arg[0-9]+]]: index -// CHECK-SAME: [[S1:%arg[0-9]+]]: index -// CHECK-SAME: [[D0:%arg[0-9]+]]: index -// CHECK-SAME: [[D1:%arg[0-9]+]]: index +// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> +// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<2x32xindex, "IAB"> +// CHECK-SAME: [[C0:%arg[0-9]+]]: index // CHECK: ktdp_lowering.construct_indirect_access_tile -// CHECK-SAME: intermediate_variables([[S0]], [[S1]], [[D0]], [[D1]]) -// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[S0]], [[S1]]{{\]}} -// CHECK-SAME: [[BASE]]{{\[}}[[C0]], [[D0]], [[D1]]{{\]}} -// CHECK-SAME: variables_space_order = #[[MAP]] -// CHECK-SAME: variables_space_set = #[[SET1]] -// CHECK-SAME: : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -// CHECK-SAME: -> !ktdp.access_tile<2x32x2x64xindex> +// CHECK-SAME: intermediate_variables([[IV0:%[a-z0-9_]+]], [[IV1:%[a-z0-9_]+]], [[IV2:%[a-z0-9_]+]], [[IV3:%[a-z0-9_]+]]) +// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[IV0]], [[IV1]]{{\]}} +// CHECK-SAME: [[BASE]][[[C0]], [[C0]] + [[IV2]], [[IV3]]] +// CHECK-SAME: variables_space_order +// CHECK-SAME: variables_space_set +// CHECK-SAME: : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -> !ktdp.access_tile<2x32x2x64xindex> func.func @roundtrip_2d_iab( %base : memref<64x2x64xf16>, %iab : memref<2x32xindex, "IAB">, - %c0 : index, - %arg5 : index, - %arg6 : index, - %arg7 : index, - %arg8 : index) { + %c0 : index) { %tile = ktdp_lowering.construct_indirect_access_tile intermediate_variables(%arg5, %arg6, %arg7, %arg8) base_ptr = %iab[%arg5, %arg6] - %base[%c0, %arg7, %arg8] - {variables_space_set = #set1, variables_space_order = #map} + %base[%c0, %c0 + %arg7, %arg8] + {variables_space_set = #set3d, variables_space_order = #map3} : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -> !ktdp.access_tile<2x32x2x64xindex> return } +// ----------------------------------------------------------------------- +// Test 2: 1-D IAB, IAB subscript is an intermediate variable. +// +// After capacity legalization: %arg5 (IAB row) absorbed into outer scf.for. +// - 3 intermediate vars: %arg6 (IAB col), %arg7, %arg8 +// - 1 captured: %c0 +// ----------------------------------------------------------------------- // CHECK-LABEL: func @roundtrip_1d_iab( -// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> -// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<32xindex, "IAB"> -// CHECK-SAME: [[C0:%arg[0-9]+]]: index -// CHECK-SAME: [[S0:%arg[0-9]+]]: index -// CHECK-SAME: [[D0:%arg[0-9]+]]: index -// CHECK-SAME: [[D1:%arg[0-9]+]]: index +// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> +// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<32xindex, "IAB"> +// CHECK-SAME: [[C0:%arg[0-9]+]]: index // CHECK: ktdp_lowering.construct_indirect_access_tile -// CHECK-SAME: intermediate_variables([[S0]], [[D0]], [[D1]]) -// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[S0]]{{\]}} -// CHECK-SAME: [[BASE]]{{\[}}[[C0]], [[D0]], [[D1]]{{\]}} -// CHECK-SAME: variables_space_order = #[[MAP]] -// CHECK-SAME: variables_space_set = #[[SET2]] -// CHECK-SAME: : memref<64x2x64xf16>, memref<32xindex, "IAB"> -// CHECK-SAME: -> !ktdp.access_tile<32x2x64xindex> +// CHECK-SAME: intermediate_variables([[IV0:%[a-z0-9_]+]], [[IV1:%[a-z0-9_]+]], [[IV2:%[a-z0-9_]+]]) +// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[IV0]]{{\]}} +// CHECK-SAME: [[BASE]][[[C0]], [[C0]] + [[IV1]], [[IV2]]] +// CHECK-SAME: variables_space_order +// CHECK-SAME: variables_space_set +// CHECK-SAME: : memref<64x2x64xf16>, memref<32xindex, "IAB"> -> !ktdp.access_tile<32x2x64xindex> func.func @roundtrip_1d_iab( %base : memref<64x2x64xf16>, %iab : memref<32xindex, "IAB">, - %c0 : index, - %arg6 : index, - %arg7 : index, - %arg8 : index) { + %c0 : index) { %tile = ktdp_lowering.construct_indirect_access_tile intermediate_variables(%arg6, %arg7, %arg8) base_ptr = %iab[%arg6] - %base[%c0, %arg7, %arg8] - {variables_space_set = #set2, variables_space_order = #map} + %base[%c0, %c0 + %arg7, %arg8] + {variables_space_set = #set2d, variables_space_order = #map2} : memref<64x2x64xf16>, memref<32xindex, "IAB"> -> !ktdp.access_tile<32x2x64xindex> return } +// ----------------------------------------------------------------------- +// Test 3: 2-D IAB where the first IAB subscript is a captured outer loop IV +// and the second is an intermediate variable (mixed). +// +// After capacity legalization: outer scf.for IV %i1 selects the IAB row +// (captured); %arg6 is the remaining intermediate variable for the IAB column. +// - 3 intermediate vars: %arg6, %arg7, %arg8 +// - 2 captured: %i1 (outer IV), %c0 +// ----------------------------------------------------------------------- // CHECK-LABEL: func @mixed_iab_subscripts( -// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> -// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<2x32xindex, "IAB"> -// CHECK-SAME: [[C0:%arg[0-9]+]]: index -// CHECK-SAME: [[I1:%arg[0-9]+]]: index -// CHECK-SAME: [[S0:%arg[0-9]+]]: index -// CHECK-SAME: [[D0:%arg[0-9]+]]: index -// CHECK-SAME: [[D1:%arg[0-9]+]]: index +// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> +// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<2x32xindex, "IAB"> +// CHECK-SAME: [[I1:%arg[0-9]+]]: index +// CHECK-SAME: [[C0:%arg[0-9]+]]: index // CHECK: ktdp_lowering.construct_indirect_access_tile -// CHECK-SAME: intermediate_variables([[S0]], [[D0]], [[D1]]) -// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[I1]], [[S0]]{{\]}} -// CHECK-SAME: [[BASE]]{{\[}}[[C0]], [[D0]], [[D1]]{{\]}} -// CHECK-SAME: variables_space_order = #[[MAP]] -// CHECK-SAME: variables_space_set = #[[SET2]] -// CHECK-SAME: : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -// CHECK-SAME: -> !ktdp.access_tile<32x2x64xindex> +// CHECK-SAME: intermediate_variables([[IV0:%[a-z0-9_]+]], [[IV1:%[a-z0-9_]+]], [[IV2:%[a-z0-9_]+]]) +// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[I1]], [[IV0]]{{\]}} +// CHECK-SAME: [[BASE]][[[C0]], [[C0]] + [[IV1]], [[IV2]]] +// CHECK-SAME: variables_space_order +// CHECK-SAME: variables_space_set +// CHECK-SAME: : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -> !ktdp.access_tile<32x2x64xindex> func.func @mixed_iab_subscripts( %base : memref<64x2x64xf16>, %iab : memref<2x32xindex, "IAB">, - %c0 : index, - %i1 : index, // outer SSA value (e.g. scf.for IV selecting IAB row) - %arg6 : index, // intermediate variable selecting IAB column - %arg7 : index, - %arg8 : index) { + %i1 : index, // captured outer loop IV selecting IAB row + %c0 : index) { %tile = ktdp_lowering.construct_indirect_access_tile intermediate_variables(%arg6, %arg7, %arg8) base_ptr = %iab[%i1, %arg6] - %base[%c0, %arg7, %arg8] - {variables_space_set = #set2, variables_space_order = #map} + %base[%c0, %c0 + %arg7, %arg8] + {variables_space_set = #set2d, variables_space_order = #map2} : memref<64x2x64xf16>, memref<2x32xindex, "IAB"> -> !ktdp.access_tile<32x2x64xindex> return } + +// ----------------------------------------------------------------------- +// Test 4: After per-entry legalization — IAB subscript is a captured scf.for +// IV (%i2); only direct-dimension intermediate vars remain. +// - 2 intermediate vars: %arg7, %arg8 +// - 2 captured: %i2 (inner loop IV), %c0 +// ----------------------------------------------------------------------- +// CHECK-LABEL: func @per_entry_legalized( +// CHECK-SAME: [[BASE:%arg[0-9]+]]: memref<64x2x64xf16> +// CHECK-SAME: [[IAB:%arg[0-9]+]]: memref<32xindex, "IAB"> +// CHECK-SAME: [[I2:%arg[0-9]+]]: index +// CHECK-SAME: [[C0:%arg[0-9]+]]: index +// CHECK: ktdp_lowering.construct_indirect_access_tile +// CHECK-SAME: intermediate_variables([[IV0:%[a-z0-9_]+]], [[IV1:%[a-z0-9_]+]]) +// CHECK-SAME: base_ptr = [[IAB]]{{\[}}[[I2]]{{\]}} +// CHECK-SAME: [[BASE]][[[C0]], [[C0]] + [[IV0]], [[IV1]]] +// CHECK-SAME: variables_space_order +// CHECK-SAME: variables_space_set +// CHECK-SAME: : memref<64x2x64xf16>, memref<32xindex, "IAB"> -> !ktdp.access_tile<2x64xindex> +func.func @per_entry_legalized( + %base : memref<64x2x64xf16>, + %iab : memref<32xindex, "IAB">, + %i2 : index, // captured inner loop IV (per-entry loop) + %c0 : index) { + %tile = ktdp_lowering.construct_indirect_access_tile + intermediate_variables(%arg7, %arg8) + base_ptr = %iab[%i2] + %base[%c0, %c0 + %arg7, %arg8] + {variables_space_set = #set2d, variables_space_order = #map2} + : memref<64x2x64xf16>, memref<32xindex, "IAB"> + -> !ktdp.access_tile<2x64xindex> + return +}