From dbb66be93e6729784c79f414634ba2da3d102963 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Mon, 20 Jul 2026 16:01:57 +0200 Subject: [PATCH] fixes slightly faster codegen --- README.md | 2 +- src/PIM/Common/IR/CompactAsmUtils.hpp | 157 ++-- src/PIM/Common/IR/CoreBlockUtils.cpp | 223 ++--- src/PIM/Common/IR/WeightUtils.cpp | 16 +- src/PIM/Common/IR/WeightUtils.hpp | 4 +- src/PIM/Compiler/CMakeLists.txt | 1 + src/PIM/Compiler/PimArtifactWriter.cpp | 61 +- src/PIM/Compiler/PimArtifactWriter.hpp | 29 + src/PIM/Compiler/PimBinaryFormat.hpp | 341 +++----- src/PIM/Compiler/PimCodeGen.cpp | 779 +++++------------- src/PIM/Compiler/PimCodeGen.hpp | 53 +- src/PIM/Compiler/PimCompilerUtils.cpp | 2 +- src/PIM/Compiler/PimCoreProgram.cpp | 209 +++++ src/PIM/Compiler/PimCoreProgram.hpp | 74 ++ src/PIM/Compiler/PimMemoryLiveness.cpp | 21 +- src/PIM/Compiler/PimMemoryLiveness.hpp | 2 - .../ONNXToSpatial/Patterns/Math/Conv.cpp | 45 +- src/PIM/Conversion/SpatialToPim/Common.cpp | 24 +- .../SpatialToPim/SpatialToPimPass.cpp | 11 +- .../OpBufferizationInterfaces.cpp | 56 +- .../Bufferization/PimBufferizationPass.cpp | 249 ++---- .../HostConstantFoldingPass.cpp | 2 +- .../MemoryCoalescing/MemoryCoalescingPass.cpp | 2 +- .../DeferredBoundaryPlanning.cpp | 68 +- .../DeferredBoundaryPlanning.hpp | 9 +- .../DeferredBoundaryRealization.cpp | 143 ++++ .../ScheduledComputeMaterialization.cpp | 16 +- 27 files changed, 1358 insertions(+), 1241 deletions(-) create mode 100644 src/PIM/Compiler/PimCoreProgram.cpp create mode 100644 src/PIM/Compiler/PimCoreProgram.hpp diff --git a/README.md b/README.md index 2f1f425..ecb7258 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,7 @@ Each validation run writes artifacts in the model workspace, for example under The compiler currently dumps dialect snapshots such as `spatial0.mlir`, `spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`, `spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`, -`pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is +`pim2_folded.mlir`, and `pim3_coalesced.mlir` when an output directory is available. To rerun the simulator manually with tracing after validation has produced a diff --git a/src/PIM/Common/IR/CompactAsmUtils.hpp b/src/PIM/Common/IR/CompactAsmUtils.hpp index 814f9d4..889189c 100644 --- a/src/PIM/Common/IR/CompactAsmUtils.hpp +++ b/src/PIM/Common/IR/CompactAsmUtils.hpp @@ -1,6 +1,7 @@ #ifndef ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP #define ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP +#include "mlir/IR/Builders.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/Value.h" #include "mlir/Support/LLVM.h" @@ -9,8 +10,11 @@ #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" +#include "llvm/ADT/Twine.h" #include "llvm/Support/LogicalResult.h" +#include + namespace onnx_mlir { namespace compact_asm { @@ -21,6 +25,59 @@ enum class ListDelimiter { Paren }; +struct NumericSuffixRange { + StringRef prefix; + unsigned first = 0; + unsigned last = 0; +}; + +inline std::optional getNumericSuffixRange(StringRef firstName, StringRef lastName) { + auto split = [](StringRef name) -> std::optional> { + size_t suffixStart = name.size(); + while (suffixStart > 0 && name[suffixStart - 1] >= '0' && name[suffixStart - 1] <= '9') + --suffixStart; + if (suffixStart == name.size()) + return std::nullopt; + + unsigned number = 0; + if (name.drop_front(suffixStart).getAsInteger(10, number)) + return std::nullopt; + return std::pair(name.take_front(suffixStart), number); + }; + + auto first = split(firstName); + auto last = split(lastName); + if (!first || !last || first->first != last->first || first->second > last->second) + return std::nullopt; + return NumericSuffixRange {first->first, first->second, last->second}; +} + +inline StringRef getInternedNumberedName(OpAsmParser& parser, StringRef prefix, unsigned number) { + return parser.getBuilder().getStringAttr((Twine(prefix) + Twine(number)).str()).getValue(); +} + +inline ParseResult parseOptionalRepeatCount(OpAsmParser& parser, int64_t& repeatCount) { + repeatCount = 1; + if (succeeded(parser.parseOptionalKeyword("x"))) { + if (parser.parseInteger(repeatCount) || repeatCount <= 0) + return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); + return success(); + } + + const char* tokenStart = parser.getCurrentLocation().getPointer(); + if (!tokenStart || tokenStart[0] != 'x' || tokenStart[1] < '0' || tokenStart[1] > '9') + return success(); + + StringRef fusedRepeat; + if (failed(parser.parseOptionalKeyword(&fusedRepeat))) + return failure(); + if (!fusedRepeat.consume_front("x") || fusedRepeat.empty() + || fusedRepeat.getAsInteger(10, repeatCount) || repeatCount <= 0) { + return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); + } + return success(); +} + inline ParseResult parseOpenDelimiter(OpAsmParser& parser, ListDelimiter delimiter) { if (delimiter == ListDelimiter::Square) return parser.parseLSquare(); @@ -59,10 +116,8 @@ inline ParseResult parseCompressedRepeatedList(OpAsmParser& parser, return failure(); int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t index = 0; index < repeatCount; ++index) entries.push_back(entry); @@ -86,10 +141,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal return failure(); int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) llvm::append_range(values, subgroup); } @@ -109,10 +162,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal return parser.emitError(parser.getCurrentLocation(), "step after 'by' must be positive"); } int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); if ((last - first) % step != 0) { return parser.emitError(parser.getCurrentLocation(), "range end must be reachable from start using the given step"); @@ -124,10 +175,8 @@ parseCompressedIntegerEntries(OpAsmParser& parser, ListDelimiter delimiter, Smal } else { int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t index = 0; index < repeatCount; ++index) values.push_back(static_cast(first)); } @@ -406,10 +455,8 @@ inline ParseResult parseCompressedTypeSequence(OpAsmParser& parser, SmallVectorI auto appendType = [&](Type type) -> ParseResult { int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t index = 0; index < repeatCount; ++index) types.push_back(type); return success(); @@ -432,18 +479,23 @@ inline ParseResult parseCompressedOperandEntryWithFirst(OpAsmParser& parser, OpAsmParser::UnresolvedOperand lastOperand; if (parser.parseOperand(lastOperand)) return failure(); - if (firstOperand.name != lastOperand.name || firstOperand.number > lastOperand.number) + if (firstOperand.name == lastOperand.name && firstOperand.number <= lastOperand.number) { + for (unsigned number = firstOperand.number; number <= lastOperand.number; ++number) + operands.push_back({firstOperand.location, firstOperand.name, number}); + return success(); + } + + auto range = getNumericSuffixRange(firstOperand.name, lastOperand.name); + if (!range || firstOperand.number != 0 || lastOperand.number != 0) return parser.emitError(parser.getCurrentLocation(), "invalid operand range"); - for (unsigned number = firstOperand.number; number <= lastOperand.number; ++number) - operands.push_back({firstOperand.location, firstOperand.name, number}); + for (unsigned number = range->first; number <= range->last; ++number) + operands.push_back({firstOperand.location, getInternedNumberedName(parser, range->prefix, number), 0}); return success(); } int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t index = 0; index < repeatCount; ++index) operands.push_back(firstOperand); return success(); @@ -560,10 +612,8 @@ inline ParseResult parseCompressedOrTupleOperandList(OpAsmParser& parser, return failure(); int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) llvm::append_range(operands, tupleOperands); @@ -574,11 +624,8 @@ inline ParseResult parseCompressedOrTupleOperandList(OpAsmParser& parser, if (parseCompressedOperandSequence(parser, tupleOperands) || parser.parseRParen()) return failure(); - repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) llvm::append_range(operands, tupleOperands); } @@ -608,10 +655,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma return failure(); int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) llvm::append_range(types, tupleTypes); @@ -622,11 +667,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma if (parseCompressedTypeSequence(parser, tupleTypes, /*allowEmpty=*/false) || parser.parseRParen()) return failure(); - repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) llvm::append_range(types, tupleTypes); } @@ -639,10 +681,8 @@ parseCompressedOrTupleTypeList(OpAsmParser& parser, ListDelimiter delimiter, Sma return failure(); int64_t repeatCount = 1; - if (succeeded(parser.parseOptionalKeyword("x"))) { - if (parser.parseInteger(repeatCount) || repeatCount <= 0) - return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive"); - } + if (parseOptionalRepeatCount(parser, repeatCount)) + return failure(); for (int64_t repeat = 0; repeat < repeatCount; ++repeat) types.push_back(type); @@ -678,13 +718,22 @@ inline ParseResult parseCompressedArgumentEntryWithFirst(OpAsmParser& parser, OpAsmParser::Argument lastArgument; if (parser.parseArgument(lastArgument)) return failure(); - if (firstArgument.ssaName.name != lastArgument.ssaName.name - || firstArgument.ssaName.number > lastArgument.ssaName.number) { - return parser.emitError(parser.getCurrentLocation(), "invalid argument range"); + if (firstArgument.ssaName.name == lastArgument.ssaName.name + && firstArgument.ssaName.number <= lastArgument.ssaName.number) { + for (unsigned number = firstArgument.ssaName.number; number <= lastArgument.ssaName.number; ++number) { + OpAsmParser::Argument argument; + argument.ssaName = {firstArgument.ssaName.location, firstArgument.ssaName.name, number}; + arguments.push_back(argument); + } + return success(); } - for (unsigned number = firstArgument.ssaName.number; number <= lastArgument.ssaName.number; ++number) { + + auto range = getNumericSuffixRange(firstArgument.ssaName.name, lastArgument.ssaName.name); + if (!range || firstArgument.ssaName.number != 0 || lastArgument.ssaName.number != 0) + return parser.emitError(parser.getCurrentLocation(), "invalid argument range"); + for (unsigned number = range->first; number <= range->last; ++number) { OpAsmParser::Argument argument; - argument.ssaName = {firstArgument.ssaName.location, firstArgument.ssaName.name, number}; + argument.ssaName = {firstArgument.ssaName.location, getInternedNumberedName(parser, range->prefix, number), 0}; arguments.push_back(argument); } return success(); diff --git a/src/PIM/Common/IR/CoreBlockUtils.cpp b/src/PIM/Common/IR/CoreBlockUtils.cpp index 9cca11a..c369411 100644 --- a/src/PIM/Common/IR/CoreBlockUtils.cpp +++ b/src/PIM/Common/IR/CoreBlockUtils.cpp @@ -3,7 +3,7 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" -#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/STLExtras.h" #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" @@ -12,10 +12,8 @@ namespace onnx_mlir { bool isCoreStaticAddressOp(mlir::Operation* op) { if (mlir::isa(op)) return true; - - if (auto selectOp = mlir::dyn_cast(op)) - return selectOp.getType().isIntOrIndex(); - - return false; + auto selectOp = mlir::dyn_cast(op); + return selectOp && selectOp.getType().isIntOrIndex(); } -mlir::LogicalResult -walkPimCoreBlock(mlir::Block& block, - const StaticValueKnowledge& initialKnowledge, - llvm::function_ref callback) { +namespace { + +enum class CoreWalkMode { ExecuteAllIterations, StructuralExtremes }; +using CoreWalkCallback = + llvm::function_ref; + +static void propagateRegionResults(mlir::ValueRange results, + mlir::Region& region, + StaticValueKnowledge& knowledge) { + if (region.empty()) + return; + auto yield = mlir::cast(region.front().getTerminator()); + for (auto [result, yielded] : llvm::zip(results, yield.getOperands())) + knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); +} + +static mlir::LogicalResult walkPimCoreBlockImpl(mlir::Block& block, + const StaticValueKnowledge& initialKnowledge, + CoreWalkMode mode, + CoreWalkCallback callback) { bool hasFailure = false; StaticValueKnowledge knowledge = initialKnowledge; + llvm::StringRef purpose = mode == CoreWalkMode::ExecuteAllIterations ? "codegen" : "verification"; for (mlir::Operation& op : block) { if (mlir::isa(op) || isCoreStaticAddressOp(&op)) continue; @@ -50,95 +62,12 @@ walkPimCoreBlock(mlir::Block& block, continue; if (auto forOp = mlir::dyn_cast(op)) { - mlir::Block& loopBody = forOp.getRegion().front(); - auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge); - auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge); + auto lower = resolveIndexValue(forOp.getLowerBound(), knowledge); + auto upper = resolveIndexValue(forOp.getUpperBound(), knowledge); auto step = resolveIndexValue(forOp.getStep(), knowledge); - if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) { - forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen"); - hasFailure = true; - continue; - } - - llvm::SmallVector iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end()); - for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) { - StaticValueKnowledge loopKnowledge = knowledge; - loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue; - for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues)) - loopKnowledge.aliases[iterArg] = iterValue; - - if (failed(walkPimCoreBlock(loopBody, loopKnowledge, callback))) - hasFailure = true; - - auto yieldOp = mlir::cast(loopBody.getTerminator()); - for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands())) - iterValues[index] = resolveLoopCarriedAlias(yieldedValue, loopKnowledge); - } - continue; - } - - if (auto ifOp = mlir::dyn_cast(op)) { - auto condition = resolveIndexValue(ifOp.getCondition(), knowledge); - if (failed(condition)) { - ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); - hasFailure = true; - continue; - } - - mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion(); - if (!selectedRegion.empty()) - if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback))) - hasFailure = true; - continue; - } - - if (auto switchOp = mlir::dyn_cast(op)) { - auto selector = resolveIndexValue(switchOp.getArg(), knowledge); - if (failed(selector)) { - switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); - hasFailure = true; - continue; - } - mlir::Region* selected = &switchOp.getDefaultRegion(); - for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions())) - if (caseValue == *selector) { - selected = &caseRegion; - break; - } - if (failed(walkPimCoreBlock(selected->front(), knowledge, callback))) - hasFailure = true; - auto yield = mlir::cast(selected->front().getTerminator()); - for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) - knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); - continue; - } - - if (failed(callback(op, knowledge))) - hasFailure = true; - } - return mlir::success(!hasFailure); -} - -mlir::LogicalResult walkPimCoreBlockStructurally( - mlir::Block& block, - const StaticValueKnowledge& initialKnowledge, - llvm::function_ref callback) { - bool hasFailure = false; - StaticValueKnowledge knowledge = initialKnowledge; - for (mlir::Operation& op : block) { - if (mlir::isa(op) || isCoreStaticAddressOp(&op)) - continue; - if (auto loadOp = mlir::dyn_cast(op); - loadOp && succeeded(resolveIndexValue(loadOp.getResult(), knowledge))) - continue; - - if (auto forOp = mlir::dyn_cast(op)) { - mlir::Block& loopBody = forOp.getRegion().front(); - auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge); - auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge); - auto step = resolveIndexValue(forOp.getStep(), knowledge); - if (failed(lowerBound) || failed(upperBound) || failed(step)) { - forOp.emitOpError("requires statically evaluable scf.for bounds for PIM verification"); + if (failed(lower) || failed(upper) || failed(step) + || (mode == CoreWalkMode::ExecuteAllIterations && *step <= 0)) { + forOp.emitOpError() << "requires statically evaluable scf.for bounds for PIM " << purpose; hasFailure = true; continue; } @@ -148,46 +77,57 @@ mlir::LogicalResult walkPimCoreBlockStructurally( continue; } - llvm::SmallVector samples; - if (*lowerBound < *upperBound) { - samples.push_back(*lowerBound); - int64_t last = *lowerBound + ((*upperBound - 1 - *lowerBound) / *step) * *step; - if (last != *lowerBound) - samples.push_back(last); - } - - for (int64_t inductionValue : samples) { + mlir::Block& body = forOp.getRegion().front(); + llvm::SmallVector iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end()); + auto visitIteration = [&](int64_t induction, bool carryValues) { StaticValueKnowledge loopKnowledge = knowledge; - loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue; - for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), forOp.getInitArgs())) - loopKnowledge.aliases[iterArg] = iterValue; + loopKnowledge.indexValues[forOp.getInductionVar()] = induction; + for (auto [index, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) + loopKnowledge.aliases[iterArg] = carryValues ? iterValues[index] : forOp.getInitArgs()[index]; + hasFailure |= failed(walkPimCoreBlockImpl(body, loopKnowledge, mode, callback)); + auto yield = mlir::cast(body.getTerminator()); + for (auto [index, yielded] : llvm::enumerate(yield.getOperands())) + iterValues[index] = resolveLoopCarriedAlias(yielded, loopKnowledge); + }; - if (failed(walkPimCoreBlockStructurally(loopBody, loopKnowledge, callback))) - hasFailure = true; + if (mode == CoreWalkMode::ExecuteAllIterations) { + for (int64_t induction = *lower; induction < *upper; induction += *step) + visitIteration(induction, true); } + else if (*lower < *upper) { + visitIteration(*lower, false); + int64_t last = *lower + ((*upper - 1 - *lower) / *step) * *step; + if (last != *lower) + visitIteration(last, false); + } + for (auto [result, value] : llvm::zip(forOp.getResults(), iterValues)) + knowledge.aliases[result] = value; continue; } if (auto ifOp = mlir::dyn_cast(op)) { - if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) { - ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification"); + auto condition = resolveIndexValue(ifOp.getCondition(), knowledge); + if (failed(condition)) { + ifOp.emitOpError() << "requires statically evaluable scf.if condition for PIM " << purpose; hasFailure = true; continue; } - - if (!ifOp.getThenRegion().empty()) - if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback))) - hasFailure = true; - if (!ifOp.getElseRegion().empty()) - if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback))) - hasFailure = true; + mlir::Region& selected = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion(); + if (mode == CoreWalkMode::ExecuteAllIterations) { + hasFailure |= !selected.empty() && failed(walkPimCoreBlockImpl(selected.front(), knowledge, mode, callback)); + } + else { + for (mlir::Region* region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()}) + hasFailure |= !region->empty() && failed(walkPimCoreBlockImpl(region->front(), knowledge, mode, callback)); + } + propagateRegionResults(ifOp.getResults(), selected, knowledge); continue; } if (auto switchOp = mlir::dyn_cast(op)) { auto selector = resolveIndexValue(switchOp.getArg(), knowledge); if (failed(selector)) { - switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification"); + switchOp.emitOpError() << "requires a statically evaluable scf.index_switch selector for PIM " << purpose; hasFailure = true; continue; } @@ -197,19 +137,34 @@ mlir::LogicalResult walkPimCoreBlockStructurally( selected = &caseRegion; break; } - for (mlir::Region& region : switchOp->getRegions()) - if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback))) - hasFailure = true; - auto yield = mlir::cast(selected->front().getTerminator()); - for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) - knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); + for (mlir::Region& region : switchOp->getRegions()) { + if (mode == CoreWalkMode::ExecuteAllIterations && ®ion != selected) + continue; + hasFailure |= failed(walkPimCoreBlockImpl(region.front(), knowledge, mode, callback)); + } + propagateRegionResults(switchOp.getResults(), *selected, knowledge); continue; } - if (failed(callback(op, knowledge))) - hasFailure = true; + hasFailure |= failed(callback(op, knowledge)); } return mlir::success(!hasFailure); } +} // namespace + +mlir::LogicalResult +walkPimCoreBlock(mlir::Block& block, + const StaticValueKnowledge& knowledge, + llvm::function_ref callback) { + return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::ExecuteAllIterations, callback); +} + +mlir::LogicalResult walkPimCoreBlockStructurally( + mlir::Block& block, + const StaticValueKnowledge& knowledge, + llvm::function_ref callback) { + return walkPimCoreBlockImpl(block, knowledge, CoreWalkMode::StructuralExtremes, callback); +} + } // namespace onnx_mlir diff --git a/src/PIM/Common/IR/WeightUtils.cpp b/src/PIM/Common/IR/WeightUtils.cpp index 395396f..ed7b775 100644 --- a/src/PIM/Common/IR/WeightUtils.cpp +++ b/src/PIM/Common/IR/WeightUtils.cpp @@ -146,17 +146,17 @@ void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight) { +std::optional resolveWeightIndex(mlir::Operation* coreLikeOp, mlir::Value weight) { weight = stripMemRefAddressingOps(weight); - if (auto coreOp = mlir::dyn_cast_or_null(weightOwner)) { + if (auto coreOp = mlir::dyn_cast_or_null(coreLikeOp)) { for (unsigned weightIndex = 0; weightIndex < coreOp.getWeights().size(); ++weightIndex) if (coreOp.getWeightArgument(weightIndex) == weight) return weightIndex; return std::nullopt; } - if (auto coreBatchOp = mlir::dyn_cast_or_null(weightOwner)) { + if (auto coreBatchOp = mlir::dyn_cast_or_null(coreLikeOp)) { for (unsigned weightIndex = 0; weightIndex < coreBatchOp.getWeights().size(); ++weightIndex) if (coreBatchOp.getWeightArgument(weightIndex) == weight) return weightIndex; @@ -167,7 +167,7 @@ std::optional resolveWeightIndex(mlir::Operation* weightOwner, mlir::V } llvm::FailureOr -resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge) { +resolveWeightView(mlir::Operation* coreLikeOp, mlir::Value weight, const StaticValueKnowledge& knowledge) { llvm::SmallVector viewOps; mlir::Value current = weight; @@ -179,7 +179,7 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static if (auto defOp = current.getDefiningOp()) { if (auto getGlobalOp = mlir::dyn_cast(defOp)) { - auto moduleOp = weightOwner ? weightOwner->getParentOfType() : mlir::ModuleOp {}; + auto moduleOp = coreLikeOp ? coreLikeOp->getParentOfType() : mlir::ModuleOp {}; auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp); if (!globalOp || !globalOp.getInitialValue()) return mlir::failure(); @@ -297,15 +297,15 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static continue; } - auto weightIndex = resolveWeightIndex(weightOwner, current); + auto weightIndex = resolveWeightIndex(coreLikeOp, current); if (!weightIndex) return mlir::failure(); - if (auto coreOp = mlir::dyn_cast_or_null(weightOwner)) { + if (auto coreOp = mlir::dyn_cast_or_null(coreLikeOp)) { current = coreOp.getWeights()[*weightIndex]; continue; } - if (auto coreBatchOp = mlir::dyn_cast_or_null(weightOwner)) { + if (auto coreBatchOp = mlir::dyn_cast_or_null(coreLikeOp)) { current = coreBatchOp.getWeights()[*weightIndex]; continue; } diff --git a/src/PIM/Common/IR/WeightUtils.hpp b/src/PIM/Common/IR/WeightUtils.hpp index 1298725..4628d6a 100644 --- a/src/PIM/Common/IR/WeightUtils.hpp +++ b/src/PIM/Common/IR/WeightUtils.hpp @@ -45,9 +45,9 @@ bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value); /// passes can identify globals that must remain weight-backed. void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref callback); -std::optional resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight); +std::optional resolveWeightIndex(mlir::Operation* coreLikeOp, mlir::Value weight); llvm::FailureOr -resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge = {}); +resolveWeightView(mlir::Operation* coreLikeOp, mlir::Value weight, const StaticValueKnowledge& knowledge = {}); template llvm::SmallVector getUsedWeightIndices(CoreLikeOpTy coreLikeOp) { diff --git a/src/PIM/Compiler/CMakeLists.txt b/src/PIM/Compiler/CMakeLists.txt index c16b63e..dd37a9d 100644 --- a/src/PIM/Compiler/CMakeLists.txt +++ b/src/PIM/Compiler/CMakeLists.txt @@ -17,6 +17,7 @@ add_pim_library(OMPimCompilerUtils PimCompilerUtils.cpp PimArtifactWriter.cpp PimCodeGen.cpp + PimCoreProgram.cpp PimMemoryLiveness.cpp PimWeightEmitter.cpp diff --git a/src/PIM/Compiler/PimArtifactWriter.cpp b/src/PIM/Compiler/PimArtifactWriter.cpp index faa9c01..fa0e247 100644 --- a/src/PIM/Compiler/PimArtifactWriter.cpp +++ b/src/PIM/Compiler/PimArtifactWriter.cpp @@ -7,7 +7,7 @@ #include #include #include -#include +#include #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp" #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp" @@ -20,6 +20,65 @@ using namespace mlir; namespace onnx_mlir { +PimInstructionWriter::PimInstructionWriter(raw_pwrite_stream& stream) +: stream(stream), buffer(kBufferSize) { + pim_binary::writeHeader(stream); +} + +PimInstructionWriter::PimInstructionWriter(raw_fd_ostream& stream) +: PimInstructionWriter(static_cast(stream)) { + fileStream = &stream; +} + +bool PimInstructionWriter::consumeStreamError() { + if (!fileStream || !fileStream->has_error()) + return false; + fileStream->clear_error(); + hasFailure = true; + return true; +} + +LogicalResult PimInstructionWriter::flushBuffer() { + if (hasFailure) + return failure(); + if (bufferedBytes != 0) { + stream.write(buffer.data(), bufferedBytes); + bufferedBytes = 0; + } + return consumeStreamError() ? failure() : success(); +} + +LogicalResult PimInstructionWriter::append(const pim_binary::InstructionRecord& record) { + if (hasFailure || finalized || count == std::numeric_limits::max()) { + hasFailure = true; + return failure(); + } + + if (bufferedBytes + pim_binary::kRecordSize > buffer.size() && failed(flushBuffer())) + return failure(); + + pim_binary::EncodedInstruction encoded = pim_binary::encodeInstructionRecord(record); + std::memcpy(buffer.data() + bufferedBytes, encoded.data(), encoded.size()); + bufferedBytes += encoded.size(); + ++count; + return success(); +} + +LogicalResult PimInstructionWriter::finalize() { + if (finalized) + return failure(); + finalized = true; + if (hasFailure || failed(flushBuffer())) + return failure(); + + stream.flush(); + if (consumeStreamError()) + return failure(); + pim_binary::patchInstructionCount(stream, count); + stream.flush(); + return consumeStreamError() ? failure() : success(); +} + OnnxMlirCompilerErrorCodes writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) { auto memoryFilePath = (outputDirPath + "/memory.bin").str(); diff --git a/src/PIM/Compiler/PimArtifactWriter.hpp b/src/PIM/Compiler/PimArtifactWriter.hpp index 0ec1e1d..d579dff 100644 --- a/src/PIM/Compiler/PimArtifactWriter.hpp +++ b/src/PIM/Compiler/PimArtifactWriter.hpp @@ -2,16 +2,45 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LogicalResult.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/JSON.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include #include "onnx-mlir/Compiler/OMCompilerTypes.h" +#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp" namespace onnx_mlir { class PimAcceleratorMemory; +class PimInstructionWriter { +public: + explicit PimInstructionWriter(llvm::raw_pwrite_stream& stream); + explicit PimInstructionWriter(llvm::raw_fd_ostream& stream); + + mlir::LogicalResult append(const pim_binary::InstructionRecord& record); + mlir::LogicalResult finalize(); + +private: + static constexpr size_t kBufferSize = 256 * 1024; + mlir::LogicalResult flushBuffer(); + bool consumeStreamError(); + + llvm::raw_pwrite_stream& stream; + llvm::raw_fd_ostream* fileStream = nullptr; + std::vector buffer; + size_t bufferedBytes = 0; + uint32_t count = 0; + bool finalized = false; + bool hasFailure = false; +}; + OnnxMlirCompilerErrorCodes writeMemoryBinary(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp, PimAcceleratorMemory& memory, diff --git a/src/PIM/Compiler/PimBinaryFormat.hpp b/src/PIM/Compiler/PimBinaryFormat.hpp index 85d0fce..889631e 100644 --- a/src/PIM/Compiler/PimBinaryFormat.hpp +++ b/src/PIM/Compiler/PimBinaryFormat.hpp @@ -1,5 +1,6 @@ #pragma once +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Endian.h" #include "llvm/Support/JSON.h" @@ -53,6 +54,14 @@ enum class Opcode : uint32_t { sync = 32, }; +inline constexpr size_t kOpcodeCount = static_cast(Opcode::sync) + 1; +inline constexpr std::array kOpcodeNames = { + "nop", "sldi", "sld", "sadd", "ssub", "smul", "saddi", "smuli", "setbw", "mvmul", "vvadd", + "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra", "vavg", "vrelu", "vtanh", "vsigm", "vsoftmax", + "vmv", "vrsu", "vrsl", "ld", "st", "lldi", "lmv", "send", "recv", "wait", "sync", +}; +static_assert(kOpcodeNames.size() == kOpcodeCount); + struct InstructionRecord { Opcode opcode = Opcode::nop; uint8_t rd = 0; @@ -64,14 +73,29 @@ struct InstructionRecord { uint8_t flags = 0; }; +using EncodedInstruction = std::array; +static_assert(kRecordSize == 20); +static_assert(sizeof(EncodedInstruction) == kRecordSize); + +inline EncodedInstruction encodeInstructionRecord(const InstructionRecord& record) { + EncodedInstruction encoded = {}; + encoded[0] = static_cast(static_cast(record.opcode)); + encoded[1] = static_cast(record.rd); + encoded[2] = static_cast(record.r1); + encoded[3] = static_cast(record.flags); + llvm::support::endian::write32le(encoded.data() + 4, static_cast(record.r2OrImm)); + llvm::support::endian::write32le(encoded.data() + 8, static_cast(record.generic1)); + llvm::support::endian::write32le(encoded.data() + 12, static_cast(record.generic2)); + llvm::support::endian::write32le(encoded.data() + 16, static_cast(record.generic3)); + return encoded; +} + inline void writeUint32LE(llvm::raw_ostream& os, uint32_t value) { std::array bytes; llvm::support::endian::write32le(bytes.data(), value); os.write(bytes.data(), bytes.size()); } -inline void writeInt32LE(llvm::raw_ostream& os, int32_t value) { writeUint32LE(os, static_cast(value)); } - inline void writeHeader(llvm::raw_ostream& os) { os.write(kMagic, sizeof(kMagic)); writeUint32LE(os, kVersion); @@ -84,17 +108,6 @@ inline void patchInstructionCount(llvm::raw_pwrite_stream& os, uint32_t instruct os.pwrite(bytes.data(), bytes.size(), kCountOffset); } -inline void writeInstructionRecord(llvm::raw_ostream& os, const InstructionRecord& record) { - os << static_cast(static_cast(record.opcode)); - os << static_cast(record.rd); - os << static_cast(record.r1); - os << static_cast(record.flags); - writeInt32LE(os, record.r2OrImm); - writeInt32LE(os, record.generic1); - writeInt32LE(os, record.generic2); - writeInt32LE(os, record.generic3); -} - inline int32_t toI32(int64_t value) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); } inline uint8_t toU8(int64_t value) { @@ -107,113 +120,64 @@ inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef return defaultValue; } +struct InstructionJsonFormat { + bool rd; + bool r1; + bool offset; + llvm::StringLiteral r2; + llvm::StringLiteral generic1; + llvm::StringLiteral generic2; + llvm::StringLiteral generic3; +}; + +inline constexpr std::array kInstructionJsonFormats = {{ + {false, false, false, "", "", "", "" }, // nop + {true, false, false, "imm", "", "", "" }, // sldi + {true, true, true, "", "", "", "" }, // sld + {true, true, false, "rs2", "", "", "" }, // sadd + {true, true, false, "rs2", "", "", "" }, // ssub + {true, true, false, "rs2", "", "", "" }, // smul + {true, true, false, "imm", "", "", "" }, // saddi + {true, true, false, "imm", "", "", "" }, // smuli + {false, false, false, "", "ibiw", "obiw", "" }, // setbw + {true, true, false, "mbiw", "relu", "group", "" }, // mvmul + {true, true, true, "rs2", "", "", "len" }, // vvadd + {true, true, true, "rs2", "", "", "len" }, // vvsub + {true, true, true, "rs2", "", "", "len" }, // vvmul + {true, true, true, "rs2", "", "", "len" }, // vvdmul + {true, true, true, "rs2", "", "", "len" }, // vvmax + {true, true, true, "rs2", "", "", "len" }, // vvsll + {true, true, true, "rs2", "", "", "len" }, // vvsra + {true, true, true, "rs2", "", "", "len" }, // vavg + {true, true, true, "", "", "", "len" }, // vrelu + {true, true, true, "", "", "", "len" }, // vtanh + {true, true, true, "", "", "", "len" }, // vsigm + {true, true, true, "", "", "", "len" }, // vsoftmax + {true, true, true, "rs2", "", "", "len" }, // vmv + {true, true, true, "rs2", "", "", "len" }, // vrsu + {true, true, true, "rs2", "", "", "len" }, // vrsl + {true, true, true, "", "", "", "size"}, // ld + {true, true, true, "", "", "", "size"}, // st + {true, false, true, "imm", "", "", "len" }, // lldi + {true, true, true, "", "", "", "len" }, // lmv + {true, false, true, "core", "", "", "size"}, // send + {true, false, true, "core", "", "", "size"}, // recv + {false, false, false, "", "", "", "" }, // wait + {false, false, false, "", "", "", "" }, // sync +}}; +static_assert(kInstructionJsonFormats.size() == kOpcodeCount); + inline Opcode opcodeFromString(llvm::StringRef opName) { - if (opName == "nop") - return Opcode::nop; - if (opName == "sldi") - return Opcode::sldi; - if (opName == "sld") - return Opcode::sld; - if (opName == "sadd") - return Opcode::sadd; - if (opName == "ssub") - return Opcode::ssub; - if (opName == "smul") - return Opcode::smul; - if (opName == "saddi") - return Opcode::saddi; - if (opName == "smuli") - return Opcode::smuli; - if (opName == "setbw") - return Opcode::setbw; - if (opName == "mvmul") - return Opcode::mvmul; - if (opName == "vvadd") - return Opcode::vvadd; - if (opName == "vvsub") - return Opcode::vvsub; - if (opName == "vvmul") - return Opcode::vvmul; - if (opName == "vvdmul") - return Opcode::vvdmul; - if (opName == "vvmax") - return Opcode::vvmax; - if (opName == "vvsll") - return Opcode::vvsll; - if (opName == "vvsra") - return Opcode::vvsra; - if (opName == "vavg") - return Opcode::vavg; - if (opName == "vrelu") - return Opcode::vrelu; - if (opName == "vtanh") - return Opcode::vtanh; - if (opName == "vsigm") - return Opcode::vsigm; - if (opName == "vsoftmax") - return Opcode::vsoftmax; - if (opName == "vmv") - return Opcode::vmv; - if (opName == "vrsu") - return Opcode::vrsu; - if (opName == "vrsl") - return Opcode::vrsl; - if (opName == "ld") - return Opcode::ld; - if (opName == "st") - return Opcode::st; - if (opName == "lldi") - return Opcode::lldi; - if (opName == "lmv") - return Opcode::lmv; - if (opName == "send") - return Opcode::send; - if (opName == "recv") - return Opcode::recv; - if (opName == "wait") - return Opcode::wait; - if (opName == "sync") - return Opcode::sync; + for (auto [index, name] : llvm::enumerate(kOpcodeNames)) + if (opName == name) + return static_cast(index); llvm_unreachable("Unsupported PIM binary opcode"); } inline llvm::StringRef opcodeToString(Opcode opcode) { - switch (opcode) { - case Opcode::nop: return "nop"; - case Opcode::sldi: return "sldi"; - case Opcode::sld: return "sld"; - case Opcode::sadd: return "sadd"; - case Opcode::ssub: return "ssub"; - case Opcode::smul: return "smul"; - case Opcode::saddi: return "saddi"; - case Opcode::smuli: return "smuli"; - case Opcode::setbw: return "setbw"; - case Opcode::mvmul: return "mvmul"; - case Opcode::vvadd: return "vvadd"; - case Opcode::vvsub: return "vvsub"; - case Opcode::vvmul: return "vvmul"; - case Opcode::vvdmul: return "vvdmul"; - case Opcode::vvmax: return "vvmax"; - case Opcode::vvsll: return "vvsll"; - case Opcode::vvsra: return "vvsra"; - case Opcode::vavg: return "vavg"; - case Opcode::vrelu: return "vrelu"; - case Opcode::vtanh: return "vtanh"; - case Opcode::vsigm: return "vsigm"; - case Opcode::vsoftmax: return "vsoftmax"; - case Opcode::vmv: return "vmv"; - case Opcode::vrsu: return "vrsu"; - case Opcode::vrsl: return "vrsl"; - case Opcode::ld: return "ld"; - case Opcode::st: return "st"; - case Opcode::lldi: return "lldi"; - case Opcode::lmv: return "lmv"; - case Opcode::send: return "send"; - case Opcode::recv: return "recv"; - case Opcode::wait: return "wait"; - case Opcode::sync: return "sync"; - } - llvm_unreachable("Unsupported PIM binary opcode"); + size_t index = static_cast(opcode); + assert(index < kOpcodeNames.size() && "Unsupported PIM binary opcode"); + return kOpcodeNames[index]; } inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruction) { @@ -221,43 +185,25 @@ inline InstructionRecord makeInstructionRecord(const llvm::json::Object& instruc std::optional opName = instruction.getString("op"); assert(opName && "Missing op field in PIM instruction"); record.opcode = opcodeFromString(*opName); - record.rd = toU8(getOptionalInt(instruction, "rd")); - record.r1 = toU8(getOptionalInt(instruction, "rs1")); - - switch (record.opcode) { - case Opcode::sldi: - case Opcode::saddi: - case Opcode::smuli: - case Opcode::lldi: record.r2OrImm = getOptionalInt(instruction, "imm"); break; - case Opcode::mvmul: - record.r2OrImm = getOptionalInt(instruction, "mbiw"); - record.generic1 = getOptionalInt(instruction, "relu"); - record.generic2 = getOptionalInt(instruction, "group"); - break; - case Opcode::setbw: - record.generic1 = getOptionalInt(instruction, "ibiw"); - record.generic2 = getOptionalInt(instruction, "obiw"); - break; - case Opcode::send: - case Opcode::recv: - record.r2OrImm = getOptionalInt(instruction, "core"); - record.generic3 = getOptionalInt(instruction, "size"); - break; - default: record.r2OrImm = getOptionalInt(instruction, "rs2"); break; - } - - if (record.opcode != Opcode::mvmul && record.opcode != Opcode::setbw) { + const auto& format = kInstructionJsonFormats[static_cast(record.opcode)]; + if (format.rd) + record.rd = toU8(getOptionalInt(instruction, "rd")); + if (format.r1) + record.r1 = toU8(getOptionalInt(instruction, "rs1")); + if (!format.r2.empty()) + record.r2OrImm = getOptionalInt(instruction, format.r2); + if (!format.generic1.empty()) + record.generic1 = getOptionalInt(instruction, format.generic1); + if (!format.generic2.empty()) + record.generic2 = getOptionalInt(instruction, format.generic2); + if (format.offset) { if (auto* offsetValue = instruction.getObject("offset")) { record.generic1 = getOptionalInt(*offsetValue, "offset_select"); record.generic2 = getOptionalInt(*offsetValue, "offset_value"); } } - - if (instruction.get("len")) - record.generic3 = getOptionalInt(instruction, "len"); - else if (instruction.get("size") && record.opcode != Opcode::send && record.opcode != Opcode::recv) - record.generic3 = getOptionalInt(instruction, "size"); - + if (!format.generic3.empty()) + record.generic3 = getOptionalInt(instruction, format.generic3); return record; } @@ -271,98 +217,21 @@ inline llvm::json::Object makeInstructionJson(const InstructionRecord& record) { offset["offset_value"] = offsetValue; instruction["offset"] = std::move(offset); }; - - switch (record.opcode) { - case Opcode::sldi: - instruction["rd"] = static_cast(record.rd); - instruction["imm"] = record.r2OrImm; - break; - case Opcode::sld: + const auto& format = kInstructionJsonFormats[static_cast(record.opcode)]; + if (format.rd) instruction["rd"] = static_cast(record.rd); + if (format.r1) instruction["rs1"] = static_cast(record.r1); + if (!format.r2.empty()) + instruction[format.r2] = record.r2OrImm; + if (!format.generic1.empty()) + instruction[format.generic1] = record.generic1; + if (!format.generic2.empty()) + instruction[format.generic2] = record.generic2; + if (format.offset) addOffset(record.generic1, record.generic2); - break; - case Opcode::sadd: - case Opcode::ssub: - case Opcode::smul: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - instruction["rs2"] = record.r2OrImm; - break; - case Opcode::saddi: - case Opcode::smuli: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - instruction["imm"] = record.r2OrImm; - break; - case Opcode::setbw: - instruction["ibiw"] = record.generic1; - instruction["obiw"] = record.generic2; - break; - case Opcode::mvmul: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - instruction["mbiw"] = record.r2OrImm; - instruction["relu"] = record.generic1; - instruction["group"] = record.generic2; - break; - case Opcode::vvadd: - case Opcode::vvsub: - case Opcode::vvmul: - case Opcode::vvdmul: - case Opcode::vvmax: - case Opcode::vvsll: - case Opcode::vvsra: - case Opcode::vavg: - case Opcode::vmv: - case Opcode::vrsu: - case Opcode::vrsl: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - instruction["rs2"] = record.r2OrImm; - addOffset(record.generic1, record.generic2); - instruction["len"] = record.generic3; - break; - case Opcode::vrelu: - case Opcode::vtanh: - case Opcode::vsigm: - case Opcode::vsoftmax: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - addOffset(record.generic1, record.generic2); - instruction["len"] = record.generic3; - break; - case Opcode::ld: - case Opcode::st: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - addOffset(record.generic1, record.generic2); - instruction["size"] = record.generic3; - break; - case Opcode::lldi: - instruction["rd"] = static_cast(record.rd); - instruction["imm"] = record.r2OrImm; - addOffset(record.generic1, record.generic2); - instruction["len"] = record.generic3; - break; - case Opcode::lmv: - instruction["rd"] = static_cast(record.rd); - instruction["rs1"] = static_cast(record.r1); - addOffset(record.generic1, record.generic2); - instruction["len"] = record.generic3; - break; - case Opcode::send: - case Opcode::recv: - instruction["rd"] = static_cast(record.rd); - instruction["core"] = record.r2OrImm; - addOffset(record.generic1, record.generic2); - instruction["size"] = record.generic3; - break; - case Opcode::wait: - case Opcode::sync: - case Opcode::nop: break; - } - + if (!format.generic3.empty()) + instruction[format.generic3] = record.generic3; return instruction; } diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 4a35aee..af343d3 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -15,7 +15,6 @@ #include "llvm/ADT/StringExtras.h" #include "llvm/Support/Debug.h" #include "llvm/Support/FileSystem.h" -#include "llvm/Support/Format.h" #include "llvm/Support/JSON.h" #include "llvm/Support/raw_ostream.h" @@ -43,6 +42,7 @@ #include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp" #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp" #include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp" #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" @@ -121,8 +121,7 @@ static FailureOr checkedAlignTo(size_t value, size_t alignment, Operatio return pim::checkedAdd(value, alignment - remainder, anchor, fieldName); } -static void printMemoryOverflowDiagnostic(mlir::Value value, - const MemoryValueKey& key, +static void printMemoryOverflowDiagnostic(const MemoryValueKey& key, size_t requestedSize, size_t currentFirstAvailableAddress, size_t alignedEndAddress) { @@ -134,10 +133,10 @@ static void printMemoryOverflowDiagnostic(mlir::Value value, if (key.lane) llvm::errs() << "Lane: " << *key.lane << "\n"; llvm::errs() << "Value: "; - value.print(llvm::errs()); + key.value.print(llvm::errs()); llvm::errs() << "\n"; - llvm::errs() << "Value type: " << value.getType() << "\n"; - if (Operation* definingOp = value.getDefiningOp()) { + llvm::errs() << "Value type: " << key.value.getType() << "\n"; + if (Operation* definingOp = key.value.getDefiningOp()) { llvm::errs() << "Defining op:\n"; definingOp->print(llvm::errs()); llvm::errs() << "\n"; @@ -146,6 +145,23 @@ static void printMemoryOverflowDiagnostic(mlir::Value value, } // namespace +size_t PimMemory::allocateAddress(size_t size, const MemoryValueKey& key) { + size_t address = firstAvailableAddress; + Operation* anchor = getDiagnosticAnchor(key.value); + auto checkedEnd = pim::checkedAdd(address, size, anchor, "local memory end"); + FailureOr checkedAlignedEnd = failure(); + if (succeeded(checkedEnd)) + checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment"); + if (address > kPimAddressLimit || failed(checkedEnd) || *checkedEnd > kPimAddressLimit + || failed(checkedAlignedEnd) || *checkedAlignedEnd > kPimAddressLimit) { + printMemoryOverflowDiagnostic( + key, size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit); + llvm_unreachable("PIM local memory allocation overflow"); + } + firstAvailableAddress = *checkedAlignedEnd; + return address; +} + MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional lane) { auto type = cast(value.getType()); assert("Only static shape is supported" && type.hasStaticShape()); @@ -170,24 +186,7 @@ void PimMemory::allocateGatheredMemory() { } void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) { - memEntry.address = firstAvailableAddress; - Operation* anchor = getDiagnosticAnchor(key.value); - auto checkedEnd = pim::checkedAdd(memEntry.address, memEntry.size, anchor, "local memory end"); - FailureOr checkedAlignedEnd = failure(); - if (succeeded(checkedEnd)) - checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment"); - bool startFits = memEntry.address <= kPimAddressLimit; - bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit; - bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit; - if (!startFits || !endFits || !alignedEndFits) { - printMemoryOverflowDiagnostic(key.value, - key, - memEntry.size, - firstAvailableAddress, - succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit); - llvm_unreachable("PIM local memory allocation overflow"); - } - firstAvailableAddress = *checkedAlignedEnd; + memEntry.address = allocateAddress(memEntry.size, key); ownedMemEntriesMap[key] = memEntry; globalMemEntriesMap[key] = memEntry; @@ -206,27 +205,8 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) { PhysicalSlotInfo slot; slot.id = nextPhysicalSlotId++; - slot.address = firstAvailableAddress; + slot.address = allocateAddress(slotSize, key); slot.size = slotSize; - - Operation* anchor = getDiagnosticAnchor(key.value); - auto checkedEnd = pim::checkedAdd(slot.address, slot.size, anchor, "local memory end"); - FailureOr checkedAlignedEnd = failure(); - if (succeeded(checkedEnd)) - checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment"); - bool startFits = slot.address <= kPimAddressLimit; - bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit; - bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit; - if (!startFits || !endFits || !alignedEndFits) { - printMemoryOverflowDiagnostic(key.value, - key, - slot.size, - firstAvailableAddress, - succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit); - llvm_unreachable("PIM local memory allocation overflow"); - } - - firstAvailableAddress = *checkedAlignedEnd; localPhysicalSlots.push_back(slot); return slot; } @@ -333,7 +313,7 @@ void PimMemory::allocateCore(Operation* op, std::optional lane) { if (pimMemoryReport != PimMemoryReportNone) { MemoryPlanArtifacts artifacts = buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport); - livenessArtifacts.textReport += artifacts.textReport; + livenessArtifacts += artifacts; } } @@ -383,23 +363,6 @@ MemoryReportRow PimMemory::getReportRow() const { return row; } -void PimMemory::remove(mlir::Value val) { - for (auto it = ownedMemEntriesMap.begin(); it != ownedMemEntriesMap.end();) - if (it->first.value == val) { - auto eraseIt = it++; - ownedMemEntriesMap.erase(eraseIt); - } - else - ++it; - for (auto it = globalMemEntriesMap.begin(); it != globalMemEntriesMap.end();) - if (it->first.value == val) { - auto eraseIt = it++; - globalMemEntriesMap.erase(eraseIt); - } - else - ++it; -} - MemEntry PimMemory::getMemEntry(const MemoryValueKey& key) const { auto iter = globalMemEntriesMap.find(key); assert("Missing memEntry for value" && iter != globalMemEntriesMap.end()); @@ -573,14 +536,6 @@ void PimAcceleratorMemory::flushReport() { fileReport.close(); } -void PimAcceleratorMemory::clean(mlir::Operation* op) { - for (auto value : op->getResults()) { - hostMem.remove(value); - for (auto& device : deviceMem) - device.second.remove(value); - } -} - size_t PimCodeGen::remapCoreId(size_t coreId) const { auto it = emittedCoreIds.find(coreId); assert(it != emittedCoreIds.end() && "Missing emitted core id remapping"); @@ -588,8 +543,8 @@ size_t PimCodeGen::remapCoreId(size_t coreId) const { } void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instruction) const { - pim_binary::writeInstructionRecord(coreBinaryStream, instruction); - ++emittedInstructionCount; + if (failed(instructionWriter.append(instruction))) + return; if (coreJsonStream) *coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ','; updateScalarRegisterCache(instruction); @@ -613,19 +568,22 @@ void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& } } -void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const { - auto registerIndex = pim::checkedU8OrCrash(registerNumber, "register number"); - auto immediateValue = pim::checkedI32OrCrash(immediate, "register immediate"); - if (scalarRegisterValues[registerIndex] == immediateValue) +void PimCodeGen::genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const { + if (scalarRegisterValues[registerNumber] == immediate) return; pim_binary::InstructionRecord instruction; instruction.opcode = pim_binary::Opcode::sldi; - instruction.rd = registerIndex; - instruction.r2OrImm = immediateValue; + instruction.rd = registerNumber; + instruction.r2OrImm = immediate; emitInstruction(instruction); } +void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const { + genSetRegisterImmediate( + pim::checkedU8OrCrash(registerNumber, "register number"), pim::checkedI32OrCrash(immediate, "register immediate")); +} + void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const { genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address")); } @@ -642,7 +600,7 @@ void PimCodeGen::setupRdRs1Rs2( genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address")); } -void PimCodeGen::emitMemCopyOp(StringRef opName, +void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode, size_t rdAddr, size_t rdOffset, size_t rs1Addr, @@ -652,7 +610,7 @@ void PimCodeGen::emitMemCopyOp(StringRef opName, setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset); pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::opcodeFromString(opName); + instruction.opcode = opcode; instruction.rd = 0; instruction.r1 = 1; instruction.generic1 = 0; @@ -661,11 +619,11 @@ void PimCodeGen::emitMemCopyOp(StringRef opName, emitInstruction(instruction); } -void PimCodeGen::emitCommunicationOp(StringRef opName, size_t bufferAddr, size_t coreId, size_t size) const { +void PimCodeGen::emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const { setupRd(bufferAddr, 0); pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::opcodeFromString(opName); + instruction.opcode = opcode; instruction.rd = 0; instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id"); instruction.generic1 = 0; @@ -692,7 +650,7 @@ void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticVa auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge); assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset) && "pim.memcp_hd offsets must be statically resolvable during codegen"); - emitMemCopyOp("ld", + emitMemCopyOp(pim_binary::Opcode::ld, addressOf(loadOp.getDeviceTarget(), knowledge), *deviceTargetOffset, addressOf(loadOp.getHostSource(), knowledge), @@ -705,7 +663,7 @@ void PimCodeGen::codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const Static auto deviceSourceOffset = indexOf(storeOp.getDeviceSourceOffset(), knowledge); assert(succeeded(hostTargetOffset) && succeeded(deviceSourceOffset) && "pim.memcp_dh offsets must be statically resolvable during codegen"); - emitMemCopyOp("st", + emitMemCopyOp(pim_binary::Opcode::st, addressOf(storeOp.getHostTarget(), knowledge), *hostTargetOffset, addressOf(storeOp.getDeviceSource(), knowledge), @@ -718,7 +676,7 @@ void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledg auto sourceOffset = indexOf(lmvOp.getSourceOffset(), knowledge); assert(succeeded(targetOffset) && succeeded(sourceOffset) && "pim.memcp offsets must be statically resolvable during codegen"); - emitMemCopyOp("lmv", + emitMemCopyOp(pim_binary::Opcode::lmv, addressOf(lmvOp.getTarget(), knowledge), *targetOffset, addressOf(lmvOp.getSource(), knowledge), @@ -730,13 +688,15 @@ void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledg void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const { auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge); assert(succeeded(sourceCoreId) && "pim.receive source core id must be statically resolvable during codegen"); - emitCommunicationOp("recv", addressOf(receiveOp.getOutputBuffer(), knowledge), *sourceCoreId, receiveOp.getSize()); + emitCommunicationOp( + pim_binary::Opcode::recv, addressOf(receiveOp.getOutputBuffer(), knowledge), *sourceCoreId, receiveOp.getSize()); } void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const { auto targetCoreId = indexOf(sendOp.getTargetCoreId(), knowledge); assert(succeeded(targetCoreId) && "pim.send target core id must be statically resolvable during codegen"); - emitCommunicationOp("send", addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize()); + emitCommunicationOp( + pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize()); } void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const { @@ -769,7 +729,8 @@ void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKno for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) { size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize; size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize; - emitMemCopyOp("lmv", outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len"); + emitMemCopyOp( + pim_binary::Opcode::lmv, outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len"); } concatOffset += inputConcatDim; @@ -786,234 +747,120 @@ void PimCodeGen::codeGenMVMLikeOp(size_t mvmId, // TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix) } -void PimCodeGen::codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vvaddOp.getOutputBuffer(), knowledge); - auto lhsAddr = addressOf(vvaddOp.getLhs(), knowledge); - auto rhsAddr = addressOf(vvaddOp.getRhs(), knowledge); - setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0); - +void PimCodeGen::emitBinaryVectorOp(pim_binary::Opcode opcode, + mlir::Value output, + mlir::Value lhs, + mlir::Value rhs, + size_t byteSize, + const StaticValueKnowledge& knowledge) const { + setupRdRs1Rs2(addressOf(output, knowledge), 0, addressOf(lhs, knowledge), 0, addressOf(rhs, knowledge), 0); pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vvadd; + instruction.opcode = opcode; instruction.rd = 0; instruction.r1 = 1; instruction.r2OrImm = 2; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vvaddOp.getLhs().getType())); + instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size"); emitInstruction(instruction); } -void PimCodeGen::codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vvsubOp.getOutputBuffer(), knowledge); - auto lhsAddr = addressOf(vvsubOp.getLhs(), knowledge); - auto rhsAddr = addressOf(vvsubOp.getRhs(), knowledge); - setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0); - +void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode, + mlir::Value output, + mlir::Value input, + size_t byteSize, + const StaticValueKnowledge& knowledge, + int32_t r2OrImm, + int32_t generic1) const { + setupRdRs1(addressOf(output, knowledge), 0, addressOf(input, knowledge), 0); pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vvsub; + instruction.opcode = opcode; instruction.rd = 0; instruction.r1 = 1; - instruction.r2OrImm = 2; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vvsubOp.getLhs().getType())); + instruction.r2OrImm = r2OrImm; + instruction.generic1 = generic1; + instruction.generic3 = pim::checkedI32OrCrash(byteSize, "vector byte size"); emitInstruction(instruction); } -void PimCodeGen::codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vvmulOp.getOutputBuffer(), knowledge); - auto lhsAddr = addressOf(vvmulOp.getLhs(), knowledge); - auto rhsAddr = addressOf(vvmulOp.getRhs(), knowledge); - setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0); +void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const { + auto srcAddr = addressOf(plan.source, knowledge); + auto dstAddr = addressOf(plan.destination, knowledge); - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vvmul; - instruction.rd = 0; - instruction.r1 = 1; - instruction.r2OrImm = 2; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vvmulOp.getLhs().getType())); - emitInstruction(instruction); -} + size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes; + int32_t maxSourceAddress = pim::checkedI32OrCrash( + pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address"); + int32_t maxDestinationAddress = pim::checkedI32OrCrash( + pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"), "transpose destination address"); + (void)maxSourceAddress; + (void)maxDestinationAddress; -void PimCodeGen::codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vvmaxOp.getOutputBuffer(), knowledge); - auto lhsAddr = addressOf(vvmaxOp.getLhs(), knowledge); - auto rhsAddr = addressOf(vvmaxOp.getRhs(), knowledge); - setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0); + pim_binary::InstructionRecord copyInstruction; + copyInstruction.opcode = pim_binary::Opcode::lmv; + copyInstruction.rd = 0; + copyInstruction.r1 = 1; - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vvmax; - instruction.rd = 0; - instruction.r1 = 1; - instruction.r2OrImm = 2; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vvmaxOp.getLhs().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vvdmulOp.getOutputBuffer(), knowledge); - auto lhsAddr = addressOf(vvdmulOp.getLhs(), knowledge); - auto rhsAddr = addressOf(vvdmulOp.getRhs(), knowledge); - setupRdRs1Rs2(outputBufferAddr, 0, lhsAddr, 0, rhsAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vvdmul; - instruction.rd = 0; - instruction.r1 = 1; - instruction.r2OrImm = 2; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vvdmulOp.getLhs().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vavgOp.getOutputBuffer(), knowledge); - auto inputAddr = addressOf(vavgOp.getInput(), knowledge); - setupRdRs1(outputBufferAddr, 0, inputAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vavg; - instruction.rd = 0; - instruction.r1 = 1; - instruction.r2OrImm = 1; - instruction.generic1 = 1; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vavgOp.getInput().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vreluOp.getOutputBuffer(), knowledge); - auto inputAddr = addressOf(vreluOp.getInput(), knowledge); - setupRdRs1(outputBufferAddr, 0, inputAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vrelu; - instruction.rd = 0; - instruction.r1 = 1; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vreluOp.getInput().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vtanhOp.getOutputBuffer(), knowledge); - auto inputAddr = addressOf(vtanhOp.getInput(), knowledge); - setupRdRs1(outputBufferAddr, 0, inputAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vtanh; - instruction.rd = 0; - instruction.r1 = 1; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vtanhOp.getInput().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vsigmOp.getOutputBuffer(), knowledge); - auto inputAddr = addressOf(vsigmOp.getInput(), knowledge); - setupRdRs1(outputBufferAddr, 0, inputAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vsigm; - instruction.rd = 0; - instruction.r1 = 1; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vsigmOp.getInput().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const { - auto outputBufferAddr = addressOf(vsoftmaxOp.getOutputBuffer(), knowledge); - auto inputAddr = addressOf(vsoftmaxOp.getInput(), knowledge); - setupRdRs1(outputBufferAddr, 0, inputAddr, 0); - - pim_binary::InstructionRecord instruction; - instruction.opcode = pim_binary::Opcode::vsoftmax; - instruction.rd = 0; - instruction.r1 = 1; - instruction.generic3 = getVectorByteSizeOrCrash(cast(vsoftmaxOp.getInput().getType())); - emitInstruction(instruction); -} - -void PimCodeGen::codeGetGlobalOp(memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const {} - -void PimCodeGen::codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const { - auto srcAddr = addressOf(transposeOp.getInput(), knowledge); - auto dstAddr = addressOf(transposeOp.getOutputBuffer(), knowledge); - - auto srcType = cast(transposeOp.getInput().getType()); - auto srcShape = srcType.getShape(); - size_t rank = srcShape.size(); - size_t elementSize = getElementTypeSizeInBytes(srcType.getElementType()); - size_t totalElements = srcType.getNumElements(); - - // Read permutation. Destination dim i corresponds to source dim perm[i]. - SmallVector perm = map_to_vector(transposeOp.getPermutation().getAsRange(), - [](auto attr) -> int64_t { return attr.getInt(); }); - - // Destination shape: dstShape[i] = srcShape[perm[i]] - SmallVector dstShape(rank); - for (size_t i = 0; i < rank; i++) - dstShape[i] = srcShape[perm[i]]; - - // Row-major strides for source and destination - SmallVector srcStrides(rank, 1); - SmallVector dstStrides(rank, 1); - for (int64_t i = rank - 2; i >= 0; i--) { - srcStrides[i] = srcStrides[i + 1] * srcShape[i + 1]; - dstStrides[i] = dstStrides[i + 1] * dstShape[i + 1]; - } - - bool storagePreserving = true; - for (size_t srcFlat = 0; srcFlat < totalElements; srcFlat++) { - SmallVector srcIdx(rank); - size_t remaining = srcFlat; - for (size_t d = 0; d < rank; d++) { - srcIdx[d] = remaining / srcStrides[d]; - remaining %= srcStrides[d]; + size_t maxRunElements = static_cast(std::numeric_limits::max()) / plan.elementBytes; + auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) { + while (runLength != 0) { + size_t chunkElements = std::min(runLength, maxRunElements); + // totalBytes was checked when the plan was compiled, so these bounded + // products cannot overflow. + size_t sourceOffset = sourceStart * plan.elementBytes; + size_t destinationOffset = destinationStart * plan.elementBytes; + size_t byteSize = chunkElements * plan.elementBytes; + assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize); + genSetRegisterImmediate(0, static_cast(dstAddr + destinationOffset)); + genSetRegisterImmediate(1, static_cast(srcAddr + sourceOffset)); + copyInstruction.generic3 = static_cast(byteSize); + emitInstruction(copyInstruction); + sourceStart += chunkElements; + destinationStart += chunkElements; + runLength -= chunkElements; } + }; - size_t dstFlat = 0; - for (size_t d = 0; d < rank; d++) - dstFlat += srcIdx[perm[d]] * dstStrides[d]; - - if (dstFlat != srcFlat) { - storagePreserving = false; - break; - } - } - - if (storagePreserving) { - emitMemCopyOp("lmv", dstAddr, 0, srcAddr, 0, totalElements * elementSize, "len"); + if (plan.storagePreserving) { + emitRun(0, 0, plan.totalElements); return; } - // Emit element-by-element copy with transposed addressing - for (size_t srcFlat = 0; srcFlat < totalElements; srcFlat++) { - // Decompose flat source index into multi-dimensional index - SmallVector srcIdx(rank); - size_t remaining = srcFlat; - for (size_t d = 0; d < rank; d++) { - srcIdx[d] = remaining / srcStrides[d]; - remaining %= srcStrides[d]; + size_t rank = plan.sourceShape.size(); + SmallVector sourceIndices(rank, 0); + size_t destinationFlat = 0; + size_t runSourceStart = 0; + size_t runDestinationStart = 0; + size_t runLength = 0; + + for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) { + if (runLength != 0 && destinationFlat != runDestinationStart + runLength) { + emitRun(runSourceStart, runDestinationStart, runLength); + runSourceStart = sourceFlat; + runDestinationStart = destinationFlat; + runLength = 0; + } + if (runLength == 0) { + runSourceStart = sourceFlat; + runDestinationStart = destinationFlat; + } + ++runLength; + + if (runLength == maxRunElements) { + emitRun(runSourceStart, runDestinationStart, runLength); + runLength = 0; } - // Compute flat destination index: dstIdx[d] = srcIdx[perm[d]] - size_t dstFlat = 0; - for (size_t d = 0; d < rank; d++) - dstFlat += srcIdx[perm[d]] * dstStrides[d]; - - emitMemCopyOp("lmv", dstAddr, dstFlat * elementSize, srcAddr, srcFlat * elementSize, elementSize, "len"); + if (sourceFlat + 1 == plan.totalElements) + break; + for (size_t sourceDim = rank; sourceDim-- > 0;) { + destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]]; + if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim]) + break; + sourceIndices[sourceDim] = 0; + destinationFlat -= plan.destinationRewinds[sourceDim]; + } } -} -size_t getMatrixSize(ShapedType matrixShape) { - if (matrixShape.getRank() != 2 && matrixShape.getRank() != 4) - assert(false && "Unsupported matrix shape"); - return std::max(matrixShape.getDimSize(0), matrixShape.getDimSize(1)); -} - -std::string getMemorySizeAsString(size_t size) { - if (size > 1024 * 1024 * 1024) - return std::to_string(size / 1024 / 1024 / 1024) + " GB"; - if (size > 1024 * 1024) - return std::to_string(size / 1024 / 1024) + " MB"; - if (size > 1024) - return std::to_string(size / 1024) + " KB"; - return std::to_string(size) + " Bytes"; + if (runLength != 0) + emitRun(runSourceStart, runDestinationStart, runLength); } static SmallVector collectTopLevelCoreLikeOps(func::FuncOp funcOp) { @@ -1093,200 +940,11 @@ public: } }; -enum class CompiledCoreOpKind : uint8_t { - Load, - Store, - Lmv, - Receive, - Send, - Concat, - Vmm, - Transpose, - VVAdd, - VVSub, - VVMul, - VVMax, - VVDMul, - VAvg, - VRelu, - VTanh, - VSigm, - VSoftmax, - GetGlobal -}; - -struct CompiledCoreNode { - enum class Kind : uint8_t { - Op, - Loop, - If, - IndexSwitch - }; - - Kind kind = Kind::Op; - Operation* op = nullptr; - CompiledCoreOpKind opKind = CompiledCoreOpKind::Load; - CompiledIndexExpr lowerBound; - CompiledIndexExpr upperBound; - CompiledIndexExpr step; - CompiledIndexExpr condition; - std::unique_ptr> loopBody; - std::unique_ptr> thenBody; - std::unique_ptr> elseBody; - llvm::SmallVector caseValues; - llvm::SmallVector>> caseBodies; - std::unique_ptr> defaultBody; -}; - -static FailureOr classifyCompiledCoreOpKind(Operation& op) { - if (isa(op)) - return CompiledCoreOpKind::Load; - if (isa(op)) - return CompiledCoreOpKind::Store; - if (isa(op)) - return CompiledCoreOpKind::Lmv; - if (isa(op)) - return CompiledCoreOpKind::Receive; - if (isa(op)) - return CompiledCoreOpKind::Send; - if (isa(op)) - return CompiledCoreOpKind::Concat; - if (isa(op)) - return CompiledCoreOpKind::Vmm; - if (isa(op)) - return CompiledCoreOpKind::Transpose; - if (isa(op)) - return CompiledCoreOpKind::VVAdd; - if (isa(op)) - return CompiledCoreOpKind::VVSub; - if (isa(op)) - return CompiledCoreOpKind::VVMul; - if (isa(op)) - return CompiledCoreOpKind::VVMax; - if (isa(op)) - return CompiledCoreOpKind::VVDMul; - if (isa(op)) - return CompiledCoreOpKind::VAvg; - if (isa(op)) - return CompiledCoreOpKind::VRelu; - if (isa(op)) - return CompiledCoreOpKind::VTanh; - if (isa(op)) - return CompiledCoreOpKind::VSigm; - if (isa(op)) - return CompiledCoreOpKind::VSoftmax; - if (isa(op)) - return CompiledCoreOpKind::GetGlobal; - return failure(); -} - -static LogicalResult -compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorImpl& plan) { - for (Operation& op : block) { - if (isa(op) || isCoreStaticAddressOp(&op)) - continue; - - if (auto loadOp = dyn_cast(op)) { - if (succeeded(compileIndexExpr(loadOp.getResult()))) - continue; - } - - if (auto forOp = dyn_cast(op)) { - auto lowerBound = compileIndexExpr(forOp.getLowerBound()); - auto upperBound = compileIndexExpr(forOp.getUpperBound()); - auto step = compileIndexExpr(forOp.getStep()); - if (failed(lowerBound) || failed(upperBound) || failed(step)) { - forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen"); - return failure(); - } - - CompiledCoreNode loopNode; - loopNode.kind = CompiledCoreNode::Kind::Loop; - loopNode.op = forOp.getOperation(); - loopNode.lowerBound = *lowerBound; - loopNode.upperBound = *upperBound; - loopNode.step = *step; - loopNode.loopBody = std::make_unique>(); - if (failed(compileCoreEmissionPlan(forOp.getRegion().front(), weightOwner, *loopNode.loopBody))) - return failure(); - plan.push_back(std::move(loopNode)); - continue; - } - - if (auto ifOp = dyn_cast(op)) { - auto condition = compileIndexExpr(ifOp.getCondition()); - if (failed(condition)) { - ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); - return failure(); - } - - CompiledCoreNode ifNode; - ifNode.kind = CompiledCoreNode::Kind::If; - ifNode.op = ifOp.getOperation(); - ifNode.condition = *condition; - ifNode.thenBody = std::make_unique>(); - if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody))) - return failure(); - ifNode.elseBody = std::make_unique>(); - if (!ifOp.getElseRegion().empty()) - if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody))) - return failure(); - plan.push_back(std::move(ifNode)); - continue; - } - - if (auto switchOp = dyn_cast(op)) { - auto selector = compileIndexExpr(switchOp.getArg()); - if (failed(selector)) { - switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); - return failure(); - } - CompiledCoreNode switchNode; - switchNode.kind = CompiledCoreNode::Kind::IndexSwitch; - switchNode.op = switchOp.getOperation(); - switchNode.condition = *selector; - llvm::append_range(switchNode.caseValues, switchOp.getCases()); - for (mlir::Region& region : switchOp.getCaseRegions()) { - auto body = std::make_unique>(); - if (failed(compileCoreEmissionPlan(region.front(), weightOwner, *body))) - return failure(); - switchNode.caseBodies.push_back(std::move(body)); - } - switchNode.defaultBody = std::make_unique>(); - if (failed(compileCoreEmissionPlan( - switchOp.getDefaultRegion().front(), weightOwner, *switchNode.defaultBody))) - return failure(); - plan.push_back(std::move(switchNode)); - continue; - } - - auto opKind = classifyCompiledCoreOpKind(op); - if (failed(opKind)) { - InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'"; - if (auto coreOp = op.getParentOfType()) - diag << " inside pim.core " << coreOp.getCoreId(); - else if (auto coreBatchOp = op.getParentOfType()) - diag << " inside pim.core_batch with laneCount " << coreBatchOp.getLaneCount(); - return failure(); - } - - CompiledCoreNode opNode; - opNode.kind = CompiledCoreNode::Kind::Op; - opNode.op = &op; - opNode.opKind = *opKind; - plan.push_back(std::move(opNode)); - } - return success(); -} - static LogicalResult executeCompiledCorePlan( const llvm::SmallVectorImpl& plan, PimCodeGen& coreCodeGen, StaticValueKnowledge& knowledge, - llvm::function_ref(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot, - size_t& processedOperations, - std::optional batchLane = std::nullopt, - std::optional batchLaneCount = std::nullopt) { + llvm::function_ref(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) { for (const CompiledCoreNode& node : plan) { if (node.kind == CompiledCoreNode::Kind::Loop) { auto lowerBound = node.lowerBound.evaluate(knowledge); @@ -1306,13 +964,7 @@ static LogicalResult executeCompiledCorePlan( for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues)) aliasBindings.bind(iterArg, iterValue); - if (failed(executeCompiledCorePlan(*node.loopBody, - coreCodeGen, - knowledge, - resolveWeightSlot, - processedOperations, - batchLane, - batchLaneCount))) + if (failed(executeCompiledCorePlan(*node.loopBody, coreCodeGen, knowledge, resolveWeightSlot))) return failure(); auto yieldOp = cast(forOp.getRegion().front().getTerminator()); @@ -1331,13 +983,7 @@ static LogicalResult executeCompiledCorePlan( } const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody; - if (selectedBody && failed(executeCompiledCorePlan(*selectedBody, - coreCodeGen, - knowledge, - resolveWeightSlot, - processedOperations, - batchLane, - batchLaneCount))) + if (selectedBody && failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot))) return failure(); continue; } @@ -1357,9 +1003,7 @@ static LogicalResult executeCompiledCorePlan( selectedRegion = &switchOp.getCaseRegions()[index]; break; } - if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, - resolveWeightSlot, processedOperations, - batchLane, batchLaneCount))) + if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, resolveWeightSlot))) return failure(); auto yield = cast(selectedRegion->front().getTerminator()); for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) @@ -1367,6 +1011,16 @@ static LogicalResult executeCompiledCorePlan( continue; } + auto emitBinary = [&](auto op, pim_binary::Opcode opcode) { + coreCodeGen.emitBinaryVectorOp(opcode, op.getOutputBuffer(), op.getLhs(), op.getRhs(), + getVectorByteSizeOrCrash(cast(op.getLhs().getType())), knowledge); + }; + auto emitUnary = [&](auto op, pim_binary::Opcode opcode, int32_t r2OrImm, int32_t generic1) { + coreCodeGen.emitUnaryVectorOp(opcode, op.getOutputBuffer(), op.getInput(), + getVectorByteSizeOrCrash(cast(op.getInput().getType())), + knowledge, r2OrImm, generic1); + }; + switch (node.opKind) { case CompiledCoreOpKind::Load: coreCodeGen.codeGenLoadOp(cast(node.op), knowledge); @@ -1385,29 +1039,33 @@ static LogicalResult executeCompiledCorePlan( return failure(); break; case CompiledCoreOpKind::Transpose: - coreCodeGen.codeGenTransposeOp(cast(node.op), knowledge); + coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break; - case CompiledCoreOpKind::VVAdd: coreCodeGen.codeGenVVAddOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VVSub: coreCodeGen.codeGenVVSubOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VVMul: coreCodeGen.codeGenVVMulOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VVMax: coreCodeGen.codeGenVVMaxOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VVDMul: coreCodeGen.codeGenVVDMulOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VAvg: coreCodeGen.codeGenVAvgOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VRelu: coreCodeGen.codeGenVReluOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VTanh: coreCodeGen.codeGenVTanhOp(cast(node.op), knowledge); break; - case CompiledCoreOpKind::VSigm: coreCodeGen.codeGenVSigmOp(cast(node.op), knowledge); break; + case CompiledCoreOpKind::VVAdd: emitBinary(cast(node.op), pim_binary::Opcode::vvadd); break; + case CompiledCoreOpKind::VVSub: emitBinary(cast(node.op), pim_binary::Opcode::vvsub); break; + case CompiledCoreOpKind::VVMul: emitBinary(cast(node.op), pim_binary::Opcode::vvmul); break; + case CompiledCoreOpKind::VVMax: emitBinary(cast(node.op), pim_binary::Opcode::vvmax); break; + case CompiledCoreOpKind::VVDMul: emitBinary(cast(node.op), pim_binary::Opcode::vvdmul); break; + case CompiledCoreOpKind::VAvg: emitUnary(cast(node.op), pim_binary::Opcode::vavg, 1, 1); break; + case CompiledCoreOpKind::VRelu: emitUnary(cast(node.op), pim_binary::Opcode::vrelu, 0, 0); break; + case CompiledCoreOpKind::VTanh: emitUnary(cast(node.op), pim_binary::Opcode::vtanh, 0, 0); break; + case CompiledCoreOpKind::VSigm: emitUnary(cast(node.op), pim_binary::Opcode::vsigm, 0, 0); break; case CompiledCoreOpKind::VSoftmax: - coreCodeGen.codeGenVSoftmaxOp(cast(node.op), knowledge); - break; - case CompiledCoreOpKind::GetGlobal: - coreCodeGen.codeGetGlobalOp(cast(node.op), knowledge); - break; + emitUnary(cast(node.op), pim_binary::Opcode::vsoftmax, 0, 0); break; } - processedOperations++; } return success(); } +static LogicalResult executeCompiledCoreProgram( + const CompiledCoreProgram& program, + PimCodeGen& coreCodeGen, + const StaticValueKnowledge& initialKnowledge, + llvm::function_ref(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot) { + StaticValueKnowledge knowledge = initialKnowledge; + return executeCompiledCorePlan(program.nodes, coreCodeGen, knowledge, resolveWeightSlot); +} + static SmallDenseMap collectMaterializedHostGlobals(ModuleOp moduleOp, func::FuncOp funcOp, const PimAcceleratorMemory& memory) { SmallDenseMap materializedHostGlobals; @@ -1444,29 +1102,6 @@ static void aliasMaterializedHostGlobals(CoreLikeOpTy coreLikeOp, }); } -/// Dispatch all operations in a core region to the appropriate code generator. -/// scf.for loops are statically unrolled via walkPimCoreBlock so that addressing is -/// fully resolved before the JSON instructions are emitted. -/// Returns the number of emitted instructions, or -1 on failure. -static int64_t codeGenCoreOps( - Block& block, - PimCodeGen& coreCodeGen, - const StaticValueKnowledge& initialKnowledge, - Operation* weightOwner, - llvm::function_ref(pim::PimVMMOp, const StaticValueKnowledge&)> resolveWeightSlot, - std::optional batchLane = std::nullopt, - std::optional batchLaneCount = std::nullopt) { - llvm::SmallVector plan; - if (failed(compileCoreEmissionPlan(block, weightOwner, plan))) - return -1; - - size_t processedOperations = 0; - StaticValueKnowledge knowledge = initialKnowledge; - auto result = executeCompiledCorePlan( - plan, coreCodeGen, knowledge, resolveWeightSlot, processedOperations, batchLane, batchLaneCount); - return failed(result) ? -1 : static_cast(processedOperations); -} - static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) { std::string outputCorePath = (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str(); @@ -1477,9 +1112,18 @@ static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath return InvalidOutputFileAccess; } - pim_binary::writeHeader(coreBinaryStream); - pim_binary::patchInstructionCount(coreBinaryStream, 0); + PimInstructionWriter instructionWriter(coreBinaryStream); + if (failed(instructionWriter.finalize())) { + coreBinaryStream.close(); + errs() << "Error while writing core file `" << outputCorePath << "`\n"; + return InvalidOutputFileAccess; + } coreBinaryStream.close(); + if (coreBinaryStream.has_error()) { + coreBinaryStream.clear_error(); + errs() << "Error while closing core file `" << outputCorePath << "`\n"; + return InvalidOutputFileAccess; + } if (!pimEmitJson.getValue()) return CompilerSuccess; @@ -1524,6 +1168,20 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: SmallVector coreLikeOps = collectTopLevelCoreLikeOps(funcOp); SmallDenseMap materializedHostGlobals = collectMaterializedHostGlobals(moduleOp, funcOp, memory); + llvm::DenseMap> compiledPrograms; + for (Operation* op : coreLikeOps) { + auto program = std::make_unique(); + if (failed(compileCoreProgram(op, *program))) + return CompilerFailure; + compiledPrograms.try_emplace(op, std::move(program)); + } + + auto getCompiledProgram = [&](Operation* op) { + auto it = compiledPrograms.find(op); + assert(it != compiledPrograms.end() && "missing compiled PIM core program"); + return it->second.get(); + }; + llvm::DenseMap emittedCoreIds; size_t nextEmittedCoreId = 0; @@ -1551,7 +1209,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: size_t originalCoreId = static_cast(coreOp.getCoreId()); CoreEmissionJob job; job.coreLikeOp = coreOp; - job.originalCoreId = originalCoreId; + job.program = getCompiledProgram(op); job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); jobs.push_back(std::move(job)); continue; @@ -1570,7 +1228,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: for (size_t originalCoreId : orderedOriginalCoreIds) { CoreEmissionJob job; job.coreLikeOp = coreBatchOp; - job.originalCoreId = originalCoreId; + job.program = getCompiledProgram(op); job.emittedCoreId = emittedCoreIds.lookup(originalCoreId); job.lanes = lanesByCoreId.lookup(originalCoreId); job.batchReportId = nextBatchReportId; @@ -1659,8 +1317,17 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: *coreJsonStream << '['; } - pim_binary::writeHeader(coreBinaryStream); - PimCodeGen coreCodeGen(jobMemory, coreBinaryStream, coreJsonStream.get(), emittedCoreIds); + PimInstructionWriter instructionWriter(coreBinaryStream); + PimCodeGen coreCodeGen(jobMemory, instructionWriter, coreJsonStream.get(), emittedCoreIds); + auto finalizeInstructions = [&]() { + bool succeeded = mlir::succeeded(instructionWriter.finalize()); + coreBinaryStream.close(); + if (coreBinaryStream.has_error()) { + coreBinaryStream.clear_error(); + succeeded = false; + } + return succeeded; + }; if (auto coreOp = dyn_cast(job.coreLikeOp)) { aliasMaterializedHostGlobals(coreOp, moduleOp, materializedHostGlobals, jobMemory); @@ -1668,9 +1335,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: deviceMemory.allocateCore(coreOp); StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp); - int64_t processedOperations = - codeGenCoreOps(coreOp.getBody().front(), coreCodeGen, knowledge, coreOp.getOperation(), resolveWeightSlot); - if (processedOperations < 0) { + if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) { + (void)finalizeInstructions(); result.status = CompilerFailure; return result; } @@ -1688,14 +1354,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: deviceMemory.allocateCore(coreBatchOp, lane); coreCodeGen.setBatchLane(lane); - int64_t processedOperations = codeGenCoreOps(coreBatchOp.getBody().front(), - coreCodeGen, - knowledge, - coreBatchOp.getOperation(), - resolveWeightSlot, - lane, - static_cast(coreBatchOp.getLaneCount())); - if (processedOperations < 0) { + if (failed(executeCompiledCoreProgram(*job.program, coreCodeGen, knowledge, resolveWeightSlot))) { + (void)finalizeInstructions(); result.status = CompilerFailure; return result; } @@ -1706,15 +1366,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: result.livenessArtifacts = deviceMemory.getLivenessArtifacts(); } - pim_binary::patchInstructionCount(coreBinaryStream, coreCodeGen.getEmittedInstructionCount()); - coreBinaryStream.close(); + if (!finalizeInstructions()) { + result.status = InvalidOutputFileAccess; + return result; + } if (coreJsonStream) { coreJsonStream->seek(coreJsonStream->tell() - 1); *coreJsonStream << ']'; coreJsonStream->close(); } - return result; }; @@ -1783,7 +1444,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup); memory.recordCoreReport(job.emittedCoreId, result.reportRow); if (livenessReportFile.is_open()) - *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts.textReport; + *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts; continue; } } @@ -1815,7 +1476,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: if (livenessReportFile.is_open()) for (size_t jobIndex : group) *livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n" - << jobResults[jobIndex].livenessArtifacts.textReport; + << jobResults[jobIndex].livenessArtifacts; } maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1; diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 078d42a..1a31ce8 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -24,6 +24,10 @@ namespace onnx_mlir { +struct CompiledCoreProgram; +struct CompiledTransposePlan; +class PimInstructionWriter; + struct MemEntry { size_t address; size_t size; @@ -35,9 +39,7 @@ struct PhysicalSlotInfo { size_t size = 0; }; -struct MemoryPlanArtifacts { - std::string textReport; -}; +using MemoryPlanArtifacts = std::string; struct MemoryValueKey { mlir::Value value; @@ -98,6 +100,7 @@ class PimMemory { size_t nextPhysicalSlotId = 0; MemEntry* gatherMemEntry(mlir::Value value, std::optional lane = std::nullopt); + size_t allocateAddress(size_t size, const MemoryValueKey& key); void allocateGatheredMemory(); void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind); PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key); @@ -110,7 +113,6 @@ public: void allocateCore(mlir::Operation* op, std::optional lane = std::nullopt); MemoryReportRow getReportRow() const; const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; } - void remove(mlir::Value val); size_t getFirstAvailableAddress() const { return firstAvailableAddress; } MemEntry getMemEntry(const MemoryValueKey& key) const; @@ -153,12 +155,11 @@ public: uint64_t totalAllocaBytes); void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; } void flushReport(); - void clean(mlir::Operation* op); }; struct CoreEmissionJob { mlir::Operation* coreLikeOp = nullptr; - size_t originalCoreId = 0; + const CompiledCoreProgram* program = nullptr; size_t emittedCoreId = 0; llvm::SmallVector lanes; std::optional batchReportId; @@ -166,11 +167,10 @@ struct CoreEmissionJob { class PimCodeGen { PimAcceleratorMemory& memory; - llvm::raw_fd_ostream& coreBinaryStream; + PimInstructionWriter& instructionWriter; llvm::raw_fd_ostream* coreJsonStream; const llvm::DenseMap& emittedCoreIds; std::optional batchLane; - mutable uint32_t emittedInstructionCount = 0; mutable std::array, 256> scalarRegisterValues = {}; size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const { @@ -181,30 +181,44 @@ class PimCodeGen { void emitInstruction(const pim_binary::InstructionRecord& instruction) const; void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const; + void genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const; void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const; void setupRd(size_t rdAddress, size_t rdOffset) const; void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const; void setupRdRs1Rs2( size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const; - void emitMemCopyOp(mlir::StringRef opName, + void emitMemCopyOp(pim_binary::Opcode opcode, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset, size_t size, mlir::StringRef sizeFieldName = "size") const; - void emitCommunicationOp(mlir::StringRef opName, size_t bufferAddr, size_t coreId, size_t size) const; + void emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const; void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const; public: + void emitBinaryVectorOp(pim_binary::Opcode opcode, + mlir::Value output, + mlir::Value lhs, + mlir::Value rhs, + size_t byteSize, + const StaticValueKnowledge& knowledge) const; + void emitUnaryVectorOp(pim_binary::Opcode opcode, + mlir::Value output, + mlir::Value input, + size_t byteSize, + const StaticValueKnowledge& knowledge, + int32_t r2OrImm = 0, + int32_t generic1 = 0) const; + PimCodeGen(PimAcceleratorMemory& memory, - llvm::raw_fd_ostream& coreBinary, + PimInstructionWriter& instructionWriter, llvm::raw_fd_ostream* coreJson, const llvm::DenseMap& emittedCoreIds) - : memory(memory), coreBinaryStream(coreBinary), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {} + : memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {} - uint32_t getEmittedInstructionCount() const { return emittedInstructionCount; } void setBatchLane(std::optional lane) { batchLane = lane; } llvm::FailureOr indexOf(mlir::Value value, const StaticValueKnowledge& knowledge) const { return memory.getIndexValue(value, knowledge); @@ -221,18 +235,7 @@ public: template void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge); - void codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowledge& knowledge) const; - void codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowledge& knowledge) const; - void codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowledge& knowledge) const; - void codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowledge& knowledge) const; - void codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKnowledge& knowledge) const; - void codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge& knowledge) const; - void codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowledge& knowledge) const; - void codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const; - void codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const; - void codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const; - void codeGetGlobalOp(mlir::memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const; - void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const; + void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const; }; OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName); diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index d5bc74d..6c9e946 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -21,7 +21,7 @@ void addPassesPim(OwningOpRef& module, verifyExplicitPimCoreCount(); if (pimOnlyCodegen) { - // Skip all the lowering passes and directly generate code for PIM. + pm.addPass(createEmitPimCodePass()); return; } diff --git a/src/PIM/Compiler/PimCoreProgram.cpp b/src/PIM/Compiler/PimCoreProgram.cpp new file mode 100644 index 0000000..612533a --- /dev/null +++ b/src/PIM/Compiler/PimCoreProgram.cpp @@ -0,0 +1,209 @@ +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "llvm/ADT/STLExtras.h" + +#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" +#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" + +using namespace llvm; +using namespace mlir; + +namespace onnx_mlir { +namespace { + +static FailureOr classifyCompiledCoreOpKind(Operation& op) { + if (isa(op)) return CompiledCoreOpKind::Load; + if (isa(op)) return CompiledCoreOpKind::Store; + if (isa(op)) return CompiledCoreOpKind::Lmv; + if (isa(op)) return CompiledCoreOpKind::Receive; + if (isa(op)) return CompiledCoreOpKind::Send; + if (isa(op)) return CompiledCoreOpKind::Concat; + if (isa(op)) return CompiledCoreOpKind::Vmm; + if (isa(op)) return CompiledCoreOpKind::Transpose; + if (isa(op)) return CompiledCoreOpKind::VVAdd; + if (isa(op)) return CompiledCoreOpKind::VVSub; + if (isa(op)) return CompiledCoreOpKind::VVMul; + if (isa(op)) return CompiledCoreOpKind::VVMax; + if (isa(op)) return CompiledCoreOpKind::VVDMul; + if (isa(op)) return CompiledCoreOpKind::VAvg; + if (isa(op)) return CompiledCoreOpKind::VRelu; + if (isa(op)) return CompiledCoreOpKind::VTanh; + if (isa(op)) return CompiledCoreOpKind::VSigm; + if (isa(op)) return CompiledCoreOpKind::VSoftmax; + return failure(); +} + +static bool isStoragePreservingTranspose(ArrayRef sourceShape, ArrayRef permutation) { + SmallVector sourceNonUnitDims; + SmallVector destinationSourceNonUnitDims; + for (auto [dim, size] : llvm::enumerate(sourceShape)) + if (size != 1) + sourceNonUnitDims.push_back(dim); + for (int64_t sourceDim : permutation) + if (sourceShape[sourceDim] != 1) + destinationSourceNonUnitDims.push_back(static_cast(sourceDim)); + return sourceNonUnitDims == destinationSourceNonUnitDims; +} + +static FailureOr compileTransposePlan(pim::PimTransposeOp transposeOp) { + auto sourceType = cast(transposeOp.getInput().getType()); + ArrayRef sourceShape = sourceType.getShape(); + size_t rank = sourceShape.size(); + CompiledTransposePlan plan; + plan.source = transposeOp.getInput(); + plan.destination = transposeOp.getOutputBuffer(); + plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType()); + auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements"); + if (failed(totalElements)) return failure(); + plan.totalElements = *totalElements; + auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size"); + if (failed(totalBytes)) return failure(); + plan.totalBytes = *totalBytes; + + SmallVector permutation = map_to_vector(transposeOp.getPermutation().getAsRange(), + [](IntegerAttr attr) { return attr.getInt(); }); + if (permutation.size() != rank) { + transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen"); + return failure(); + } + + SmallVector destinationShape(rank); + plan.destinationStrides.assign(rank, 1); + plan.destinationDimensionForSource.assign(rank, 0); + plan.destinationRewinds.assign(rank, 0); + SmallVector seenSourceDimensions(rank, false); + for (size_t dim = 0; dim < rank; ++dim) { + auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension"); + if (failed(size)) return failure(); + plan.sourceShape.push_back(*size); + } + for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) { + if (sourceDim < 0 || static_cast(sourceDim) >= rank || seenSourceDimensions[sourceDim]) { + transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once"); + return failure(); + } + seenSourceDimensions[sourceDim] = true; + destinationShape[destinationDim] = plan.sourceShape[sourceDim]; + plan.destinationDimensionForSource[sourceDim] = destinationDim; + } + for (size_t dim = rank; dim > 1; --dim) { + auto stride = pim::checkedMul( + plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride"); + if (failed(stride)) return failure(); + plan.destinationStrides[dim - 2] = *stride; + } + for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) { + auto rewind = pim::checkedMul(plan.sourceShape[sourceDim], + plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]], + transposeOp, + "transpose destination rewind"); + if (failed(rewind)) return failure(); + plan.destinationRewinds[sourceDim] = *rewind; + } + plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation); + return plan; +} + +static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl& plan) { + for (Operation& op : block) { + if (isa(op) || isCoreStaticAddressOp(&op)) + continue; + if (auto loadOp = dyn_cast(op); loadOp && succeeded(compileIndexExpr(loadOp.getResult()))) + continue; + + if (auto forOp = dyn_cast(op)) { + auto lower = compileIndexExpr(forOp.getLowerBound()); + auto upper = compileIndexExpr(forOp.getUpperBound()); + auto step = compileIndexExpr(forOp.getStep()); + if (failed(lower) || failed(upper) || failed(step)) { + forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen"); + return failure(); + } + CompiledCoreNode node; + node.kind = CompiledCoreNode::Kind::Loop; + node.op = forOp; + node.lowerBound = *lower; + node.upperBound = *upper; + node.step = *step; + node.loopBody = std::make_unique>(); + if (failed(compileCoreEmissionPlan(forOp.getRegion().front(), *node.loopBody))) return failure(); + plan.push_back(std::move(node)); + continue; + } + if (auto ifOp = dyn_cast(op)) { + auto condition = compileIndexExpr(ifOp.getCondition()); + if (failed(condition)) { + ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); + return failure(); + } + CompiledCoreNode node; + node.kind = CompiledCoreNode::Kind::If; + node.op = ifOp; + node.condition = *condition; + node.thenBody = std::make_unique>(); + node.elseBody = std::make_unique>(); + if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), *node.thenBody))) return failure(); + if (!ifOp.getElseRegion().empty() + && failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), *node.elseBody))) + return failure(); + plan.push_back(std::move(node)); + continue; + } + if (auto switchOp = dyn_cast(op)) { + auto selector = compileIndexExpr(switchOp.getArg()); + if (failed(selector)) { + switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); + return failure(); + } + CompiledCoreNode node; + node.kind = CompiledCoreNode::Kind::IndexSwitch; + node.op = switchOp; + node.condition = *selector; + llvm::append_range(node.caseValues, switchOp.getCases()); + for (Region& region : switchOp.getCaseRegions()) { + auto body = std::make_unique>(); + if (failed(compileCoreEmissionPlan(region.front(), *body))) return failure(); + node.caseBodies.push_back(std::move(body)); + } + node.defaultBody = std::make_unique>(); + if (failed(compileCoreEmissionPlan(switchOp.getDefaultRegion().front(), *node.defaultBody))) return failure(); + plan.push_back(std::move(node)); + continue; + } + + auto opKind = classifyCompiledCoreOpKind(op); + if (failed(opKind)) { + InFlightDiagnostic diagnostic = op.emitError() << "unsupported codegen for op '" << op.getName() << "'"; + if (auto coreOp = op.getParentOfType()) + diagnostic << " inside pim.core " << coreOp.getCoreId(); + else if (auto batchOp = op.getParentOfType()) + diagnostic << " inside pim.core_batch with laneCount " << batchOp.getLaneCount(); + return failure(); + } + CompiledCoreNode node; + node.op = &op; + node.opKind = *opKind; + if (*opKind == CompiledCoreOpKind::Transpose) { + auto transposePlan = compileTransposePlan(cast(op)); + if (failed(transposePlan)) return failure(); + node.transposePlan = *transposePlan; + } + plan.push_back(std::move(node)); + } + return success(); +} + +} // namespace + +LogicalResult compileCoreProgram(Operation* coreLikeOp, CompiledCoreProgram& program) { + Block& block = isa(coreLikeOp) ? cast(coreLikeOp).getBody().front() + : cast(coreLikeOp).getBody().front(); + return compileCoreEmissionPlan(block, program.nodes); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCoreProgram.hpp b/src/PIM/Compiler/PimCoreProgram.hpp new file mode 100644 index 0000000..312714b --- /dev/null +++ b/src/PIM/Compiler/PimCoreProgram.hpp @@ -0,0 +1,74 @@ +#pragma once + +#include "mlir/IR/Operation.h" +#include "mlir/Support/LogicalResult.h" + +#include "llvm/ADT/SmallVector.h" + +#include +#include + +#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp" + +namespace onnx_mlir { + +struct CompiledTransposePlan { + mlir::Value source; + mlir::Value destination; + size_t elementBytes = 0; + size_t totalElements = 0; + size_t totalBytes = 0; + llvm::SmallVector sourceShape; + llvm::SmallVector destinationStrides; + llvm::SmallVector destinationDimensionForSource; + llvm::SmallVector destinationRewinds; + bool storagePreserving = false; +}; + +enum class CompiledCoreOpKind : uint8_t { + Load, + Store, + Lmv, + Receive, + Send, + Concat, + Vmm, + Transpose, + VVAdd, + VVSub, + VVMul, + VVMax, + VVDMul, + VAvg, + VRelu, + VTanh, + VSigm, + VSoftmax +}; + +struct CompiledCoreNode { + enum class Kind : uint8_t { Op, Loop, If, IndexSwitch }; + + Kind kind = Kind::Op; + mlir::Operation* op = nullptr; + CompiledCoreOpKind opKind = CompiledCoreOpKind::Load; + CompiledIndexExpr lowerBound; + CompiledIndexExpr upperBound; + CompiledIndexExpr step; + CompiledIndexExpr condition; + std::unique_ptr> loopBody; + std::unique_ptr> thenBody; + std::unique_ptr> elseBody; + llvm::SmallVector caseValues; + llvm::SmallVector>> caseBodies; + std::unique_ptr> defaultBody; + std::optional transposePlan; +}; + +struct CompiledCoreProgram { + llvm::SmallVector nodes; +}; + +mlir::LogicalResult compileCoreProgram(mlir::Operation* coreLikeOp, CompiledCoreProgram& program); + +} // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimMemoryLiveness.cpp b/src/PIM/Compiler/PimMemoryLiveness.cpp index 5deeadb..a53a998 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.cpp +++ b/src/PIM/Compiler/PimMemoryLiveness.cpp @@ -42,8 +42,6 @@ static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional& aliases, struct OrderedTouchRange { uint64_t start = 0; uint64_t end = 0; - Operation* startOp = nullptr; - Operation* endOp = nullptr; bool escapedLoop = false; }; static OrderedTouchRange getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) { - OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), user, user, false}; + OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), false}; for (Operation* current = user; current; current = current->getParentOp()) { auto forOp = dyn_cast(current); if (!forOp || isWithin(definingValue, &forOp.getRegion())) continue; range.start = std::min(range.start, ordering.position.lookup(forOp)); range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp)); - range.startOp = forOp; - range.endOp = forOp; range.escapedLoop = true; } return range; @@ -247,8 +241,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, MemoryTouchInterval interval; interval.start = ordering.position.lookup(allocOp); interval.end = interval.start; - interval.startOp = allocOp; - interval.endOp = allocOp; auto recordAlias = [&](mlir::Value value) { if (includeAliasDescriptions) appendAliasDescription(interval.aliasesFollowed, value); @@ -350,18 +342,14 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, if (!interval.hasRuntimeUse) { interval.start = range.start; interval.end = range.end; - interval.startOp = range.startOp; - interval.endOp = range.endOp; interval.hasRuntimeUse = true; } else { if (range.start < interval.start) { interval.start = range.start; - interval.startOp = range.startOp; } if (range.end > interval.end) { interval.end = range.end; - interval.endOp = range.endOp; } } continue; @@ -380,8 +368,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, interval.endUsedFallback = true; interval.start = ordering.position.lookup(allocOp); interval.end = fallbackEnd; - interval.startOp = allocOp; - interval.endOp = allocOp->getParentOp(); interval.firstTouchPosition = interval.start; interval.lastTouchPosition = interval.end; addFallbackReason(interval.fallbackReason, "no runtime memory touch"); @@ -390,7 +376,6 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, if (interval.endUsedFallback) { interval.end = std::max(interval.end, fallbackEnd); - interval.endOp = allocOp->getParentOp(); } return interval; @@ -447,8 +432,6 @@ SmallVector onnx_mlir::buildLocalAllocIntervals(Operation interval.start = touchInterval.start; interval.end = touchInterval.end; interval.size = *checkedSize; - interval.startOp = touchInterval.startOp; - interval.endOp = touchInterval.endOp; interval.firstTouchOp = touchInterval.firstTouchOp; interval.lastTouchOp = touchInterval.lastTouchOp; interval.firstTouchPosition = touchInterval.firstTouchPosition; @@ -574,7 +557,7 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp, double savedPercent = totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast(savedBytes) / static_cast(totalLogicalBytes); - raw_string_ostream os(artifacts.textReport); + raw_string_ostream os(artifacts); os << "=== PIM Memory Liveness Report ===\n"; os << "Op: " << coreLikeOp->getName() << "\n"; if (lane) diff --git a/src/PIM/Compiler/PimMemoryLiveness.hpp b/src/PIM/Compiler/PimMemoryLiveness.hpp index 99c8da7..5a25174 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.hpp +++ b/src/PIM/Compiler/PimMemoryLiveness.hpp @@ -21,8 +21,6 @@ struct LocalAllocInterval { uint64_t start = 0; uint64_t end = 0; size_t size = 0; - mlir::Operation* startOp = nullptr; - mlir::Operation* endOp = nullptr; mlir::Operation* firstTouchOp = nullptr; mlir::Operation* lastTouchOp = nullptr; uint64_t firstTouchPosition = 0; diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 74e4011..09817f2 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -2180,39 +2180,55 @@ static Value createIm2colRows(const ConvLoweringState& state, auto elemType = preparedInput.type.getElementType(); auto packedRowType = RankedTensorType::get( {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding()); - auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType)); - Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType); + auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType); + auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType); + bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0; + SmallVector im2colInputs {preparedInput.value}; auto im2colComputeOp = createSpatComputeBatch( rewriter, loc, TypeRange {plan.gemmInputRowsType}, plan.packedNumRows, {}, - ValueRange {preparedInput.value, zeroRow}, + im2colInputs, [&](detail::SpatComputeBatchBodyArgs args) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels); - Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches); Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp); - Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart); - Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack); - Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack); - auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType); - auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType); - + Value lanePatches = cPack; + if (hasPartialLane) { + Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches); + Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart); + Value isPartial = arith::CmpIOp::create( + rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack); + lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack); + } + Value rowInit = tensor::EmptyOp::create(rewriter, loc, packedRowType.getShape(), elemType); + if (hasPartialLane) { + auto zeroAttr = cast(rewriter.getZeroAttr(elemType)); + rowInit = linalg::MapOp::create( + rewriter, loc, ValueRange {}, rowInit, + [&](OpBuilder& builder, Location nestedLoc, ValueRange) { + Value zero = arith::ConstantOp::create(builder, nestedLoc, zeroAttr); + linalg::YieldOp::create(builder, nestedLoc, zero); + }).getResult().front(); + } auto rowLoop = buildNormalizedScfFor( rewriter, loc, c0, lanePatches, c1, - ValueRange {args.inputs[1]}, + ValueRange {rowInit}, [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex); - Value batchIndex = - affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp); + Value batchIndex = state.batchSize == 1 + ? c0 + : affineAddFloorDivConst( + rewriter, nestedLoc, patchIndex, plan.chunkStart, + plan.numPatchesPerBatch, anchorOp); Value batchPatchIndex = affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp); Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp); @@ -2249,7 +2265,8 @@ static Value createIm2colRows(const ConvLoweringState& state, }); if (failed(rowLoop)) return failure(); - publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane); + Value row = rowLoop->results.front(); + publishGraphBatchPhysicalFragment(rewriter, loc, row, args.outputs.front(), args.lane); return success(); }); diff --git a/src/PIM/Conversion/SpatialToPim/Common.cpp b/src/PIM/Conversion/SpatialToPim/Common.cpp index d34912b..451d45b 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.cpp +++ b/src/PIM/Conversion/SpatialToPim/Common.cpp @@ -1,6 +1,7 @@ #include "mlir/IR/ValueRange.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "llvm/ADT/STLExtras.h" @@ -59,6 +60,23 @@ bool hasLaterUserInBlock(mlir::Value value, Operation* operation) { return false; } +static bool isTensorView(mlir::Value value) { + return isa_and_nonnull(value.getDefiningOp()); +} + +static bool isLoopCarriedOutput(mlir::Value operand, Operation* operation) { + auto argument = dyn_cast(operand); + if (!argument || argument.getArgNumber() == 0 || operation->getBlock() != argument.getOwner()) + return false; + auto loop = dyn_cast_or_null(argument.getOwner()->getParentOp()); + return loop && cast(loop.getBody()->getTerminator()) + .getOperand(argument.getArgNumber() - 1) == operation->getResult(0); +} + mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) { assert("Only support operations with a single result" && operation->getNumResults() == 1); mlir::Value result = operation->getResult(0); @@ -67,7 +85,11 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op SmallVector operands = getOpOperandsSortedByUses(operation); auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) { - return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation); + return operand.getType() == resultType + && (!isa(operand) || isLoopCarriedOutput(operand, operation)) + && !operand.getDefiningOp() + && !isTensorView(operand) + && !hasLaterUserInBlock(operand, operation); }); auto bestOperand = validOperands.begin(); diff --git a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp index a9aa223..44c0ee6 100644 --- a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp +++ b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp @@ -2,6 +2,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SCF/Utils/Utils.h" @@ -112,7 +113,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { memref::MemRefDialect, scf::SCFDialect, BuiltinDialect>(); - target.addLegalOp(); @@ -186,7 +189,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { memref::MemRefDialect, scf::SCFDialect, BuiltinDialect>(); - coreBodyTarget.addLegalOp(); @@ -234,7 +239,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { memref::MemRefDialect, scf::SCFDialect, BuiltinDialect>(); - communicationTarget.addLegalOp(); + communicationTarget.addLegalOp(); communicationTarget.addIllegalOp +struct CopyOpInterface +: DstBufferizableOpInterfaceExternalModel { + bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { + return cast(op).isDpsInput(&opOperand); + } + + bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { + return cast(op).isDpsInit(&opOperand); + } + + AliasingValueList getAliasingValues(Operation* op, + OpOperand& opOperand, + const AnalysisState& state) const { + auto dstOp = cast(op); + if (dstOp.isDpsInit(&opOperand)) + return {{dstOp.getTiedOpResult(&opOperand), BufferRelation::Equivalent}}; + return {}; + } + + AliasingOpOperandList getAliasingOpOperands(Operation* op, + Value value, + const AnalysisState& state) const { + auto result = dyn_cast(value); + if (!result || result.getDefiningOp() != op) + return {}; + return {{cast(op).getTiedOpOperand(result), BufferRelation::Equivalent}}; + } + + bool mustBufferizeInPlace(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { + return isa(opOperand.get().getType()); + } + + bool isWritable(Operation* op, Value value, const AnalysisState& state) const { return isa(value); } + + bool isNotConflicting(Operation* op, + OpOperand* read, + OpOperand* write, + const AnalysisState& state) const { + if (read->getOwner() != op || write->getOwner() != op) + return false; + auto dstOp = cast(op); + return dstOp.isDpsInput(read) && dstOp.isDpsInit(write); + } +}; + struct MemCopyHostToDevOpInterface -: DstBufferizableOpInterfaceExternalModel { +: CopyOpInterface { LogicalResult bufferize(Operation* op, RewriterBase& rewriter, const BufferizationOptions& options, @@ -68,7 +114,7 @@ struct MemCopyHostToDevOpInterface }; struct MemCopyDevToHostOpInterface -: DstBufferizableOpInterfaceExternalModel { +: CopyOpInterface { LogicalResult bufferize(Operation* op, RewriterBase& rewriter, const BufferizationOptions& options, @@ -99,11 +145,7 @@ struct MemCopyDevToHostOpInterface } }; -struct MemCopyOpInterface : DstBufferizableOpInterfaceExternalModel { - bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { - return !cast(op).isDpsInit(&opOperand); - } - +struct MemCopyOpInterface : CopyOpInterface { LogicalResult bufferize(Operation* op, RewriterBase& rewriter, const BufferizationOptions& options, diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp index b807f03..cc98bde 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp @@ -1,5 +1,4 @@ #include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h" #include "mlir/Dialect/Bufferization/Transforms/Transforms.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -7,16 +6,15 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/Dominance.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" -#include "mlir/Rewrite/PatternApplicator.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "llvm/ADT/SmallPtrSet.h" -#include "llvm/ADT/SmallSet.h" #include "llvm/Support/Casting.h" #include "Common/PimCommon.hpp" +#include "Common/Support/Diagnostics.hpp" #include "Compiler/PimCodeGen.hpp" #include "Dialect/Pim/PimOps.hpp" #include "Dialect/Pim/Transforms/Bufferization/Common.hpp" @@ -57,7 +55,10 @@ static StaticValueKnowledge seedCoreBatchKnowledge(pim::PimCoreBatchOp coreBatch } static LogicalResult -lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const StaticValueKnowledge& knowledge) { +lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, + Value zeroOffset, + PatternRewriter& rewriter, + const StaticValueKnowledge& knowledge) { if (!copyOp->getParentOfType() && !copyOp->getParentOfType()) return failure(); @@ -68,7 +69,6 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const if (sourceType.getElementType() != targetType.getElementType()) return failure(); - Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0); auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource()); if (failed(sizeAttr)) return failure(); @@ -121,58 +121,35 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const return success(); } -static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp, - const StaticValueKnowledge& knowledge, - bool emitDiagnostic) { - bool sourceIsHost = isHostBackedPimAddress(copyOp.getHostSource(), knowledge); - bool targetIsHost = isHostBackedPimAddress(copyOp.getDeviceTarget(), knowledge); - bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getHostSource(), knowledge); - bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceTarget(), knowledge); - if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice) { - if (emitDiagnostic) - copyOp.emitOpError() << "pim.memcp_hd requires a host-backed source and a device-local target: source=" - << copyOp.getHostSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice - << ", target=" << copyOp.getDeviceTarget() << " host=" << targetIsHost - << " device=" << targetIsDevice; - return failure(); - } - return success(); -} +enum class ExpectedPimCopyDirection { HostToDevice, DeviceToHost, DeviceToDevice }; -static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp, - const StaticValueKnowledge& knowledge, - bool emitDiagnostic) { - bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge); - bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge); - bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge); - bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge); - if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice) { - if (emitDiagnostic) - copyOp.emitOpError() << "pim.memcp_dh requires a device-local source and a host-backed target: source=" - << copyOp.getDeviceSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice - << ", target=" << copyOp.getHostTarget() << " host=" << targetIsHost - << " device=" << targetIsDevice; - return failure(); - } - return success(); -} - -static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp, - const StaticValueKnowledge& knowledge, - bool emitDiagnostic) { - bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge); - bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge); - bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge); - bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge); - if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost) { - if (emitDiagnostic) - copyOp.emitOpError() << "pim.memcp requires device-local source and target operands: source=" - << copyOp.getSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice - << ", target=" << copyOp.getTarget() << " host=" << targetIsHost - << " device=" << targetIsDevice; - return failure(); - } - return success(); +static LogicalResult verifyPimCopyEndpoints(Operation* copy, + Value source, + Value target, + ExpectedPimCopyDirection direction, + const StaticValueKnowledge& knowledge, + bool emitDiagnostic) { + struct ExpectedEndpoints { + bool sourceHost, sourceDevice, targetHost, targetDevice; + StringLiteral description; + }; + static constexpr ExpectedEndpoints expected[] = { + {true, false, false, true, "a host-backed source and a device-local target"}, + {false, true, true, false, "a device-local source and a host-backed target"}, + {false, true, false, true, "device-local source and target operands"}, + }; + const auto& endpoints = expected[static_cast(direction)]; + bool sourceHost = isHostBackedPimAddress(source, knowledge); + bool sourceDevice = isDeviceLocalPimAddress(source, knowledge); + bool targetHost = isHostBackedPimAddress(target, knowledge); + bool targetDevice = isDeviceLocalPimAddress(target, knowledge); + bool valid = sourceHost == endpoints.sourceHost && sourceDevice == endpoints.sourceDevice + && targetHost == endpoints.targetHost && targetDevice == endpoints.targetDevice; + if (!valid && emitDiagnostic) + copy->emitOpError() << "requires " << endpoints.description << ": source=" << source << " host=" << sourceHost + << " device=" << sourceDevice << ", target=" << target << " host=" << targetHost + << " device=" << targetDevice; + return success(valid); } struct PimBufferizationPass : PassWrapper> { @@ -191,14 +168,6 @@ private: LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const; }; -static LogicalResult applyPatternsOnce(Operation* op, PatternApplicator& applicator, PatternRewriter& rewriter) { - if (!op || !op->getBlock()) - return failure(); - - rewriter.setInsertionPoint(op); - return applicator.matchAndRewrite(op, rewriter); -} - static void materializeWritableConstantDestinations(func::FuncOp funcOp) { SmallVector constantBackedRoots; llvm::SmallPtrSet seenRoots; @@ -237,69 +206,44 @@ static void materializeWritableConstantDestinations(func::FuncOp funcOp) { } } -static LogicalResult verifyConflictFreePimCoreWrites( - func::FuncOp funcOp, const bufferization::OneShotBufferizationOptions& options) { - bufferization::AnalysisState analysisState(options); - DominanceInfo dominance(funcOp); - size_t violationCount = 0; +static LogicalResult verifyPimCoresNeedNoTensorCopies( + ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) { + static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc"; + OwningOpRef clone = module.clone(); + clone->walk([&](bufferization::AllocTensorOp alloc) { + if (alloc->getParentOfType() + || alloc->getParentOfType()) + alloc->setAttr(kExistingAlloc, UnitAttr::get(module.getContext())); + }); - auto verifyCore = [&](Operation* coreOp) { - coreOp->walk([&](Operation* writeOp) { - for (OpOperand& write : writeOp->getOpOperands()) { - if (!isa(write.get().getType()) || !analysisState.bufferizesToMemoryWrite(write)) - continue; + auto options = baseOptions; + options.bufferizeFunctionBoundaries = false; + options.opFilter.allowOperation([](Operation* op) { + return isa(op) + || op->getParentOfType() + || op->getParentOfType(); + }); - SmallVector worklist {write.get()}; - llvm::SmallDenseSet visited; - bool hasConflict = false; - Value conflictingAlias; - Operation* conflictingUse = nullptr; - while (!worklist.empty() && !hasConflict) { - Value alias = worklist.pop_back_val(); - if (!visited.insert(alias).second) - continue; + bufferization::BufferizationState state; + if (failed(bufferization::insertTensorCopies(*clone, options, state))) { + module.emitError("official one-shot analysis failed while verifying PIM core copy freedom"); + return failure(); + } - for (OpOperand& use : alias.getUses()) { - bool usePrecedesWrite = false; - for (Operation* ancestor = use.getOwner(); ancestor; ancestor = ancestor->getParentOp()) - if (ancestor == writeOp || dominance.properlyDominates(ancestor, writeOp)) { - usePrecedesWrite = true; - break; - } - if (usePrecedesWrite || analysisState.insideMutuallyExclusiveRegions(use.getOwner(), writeOp)) - continue; - hasConflict = true; - conflictingAlias = alias; - conflictingUse = use.getOwner(); - break; - } - - if (alias.getDefiningOp() - && llvm::all_of(alias.getUses(), [&](OpOperand& use) { - return use.getOwner() == writeOp; - })) - continue; - if (isa(alias)) - for (bufferization::AliasingOpOperand tied : analysisState.getAliasingOpOperands(alias).getAliases()) - worklist.push_back(tied.opOperand->get()); - } - - if (!hasConflict) - continue; - if (violationCount++ == 0) - writeOp->emitOpError() << "PIM core tensor write may modify an alias used later: operand #" - << write.getOperandNumber() << " (" << write.get() << "), alias=" - << conflictingAlias << ", later use=" << *conflictingUse; - } + CappedDiagnosticReporter diagnostics; + clone->walk([&](bufferization::AllocTensorOp alloc) { + if (alloc->hasAttr(kExistingAlloc) + || (!alloc->getParentOfType() + && !alloc->getParentOfType())) + return; + Operation* requiredBy = alloc->getUsers().empty() + ? alloc.getOperation() : *alloc->getUsers().begin(); + diagnostics.report(requiredBy, [](Operation* op) { + op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core"); }); - }; - - funcOp.walk([&](pim::PimCoreOp coreOp) { verifyCore(coreOp); }); - funcOp.walk([&](pim::PimCoreBatchOp coreOp) { verifyCore(coreOp); }); - if (violationCount != 0) - funcOp.emitError() << "found " << violationCount - << " non-linear PIM tensor write(s); the first is reported above"; - return success(violationCount == 0); + }); + diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies"); + return success(!diagnostics.hasFailure()); } } // namespace @@ -314,7 +258,7 @@ void PimBufferizationPass::runOnOperation() { options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap); materializeWritableConstantDestinations(funcOp); - if (failed(verifyConflictFreePimCoreWrites(funcOp, options))) { + if (failed(verifyPimCoresNeedNoTensorCopies(moduleOp, options))) { signalPassFailure(); return; } @@ -325,14 +269,8 @@ void PimBufferizationPass::runOnOperation() { || op->getParentOfType(); }); bufferization::BufferizationState state; - if (failed(bufferization::insertTensorCopies( - moduleOp, hostOptions, state))) { - moduleOp.emitError("Failed to bufferize PIM and Spatial ops"); - signalPassFailure(); - return; - } - if (failed(bufferization::bufferizeModuleOp( - moduleOp, options, state))) { + if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state)) + || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) { moduleOp.emitError("Failed to bufferize PIM and Spatial ops"); signalPassFailure(); return; @@ -369,15 +307,16 @@ void PimBufferizationPass::runOnOperation() { if (auto copyOp = dyn_cast(&op)) addCopyOp(copyOp, opKnowledge); return success(); - }); + }); } }); bool hasFailed = false; + Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0); for (const MemRefCopyWorkItem& workItem : copyWorklist) { memref::CopyOp copyOp = workItem.copyOp; rewriter.setInsertionPoint(copyOp); - if (failed(lowerMemRefCopyToPimCopy(copyOp, rewriter, workItem.knowledge))) + if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge))) hasFailed = true; } if (hasFailed) { @@ -387,32 +326,10 @@ void PimBufferizationPass::runOnOperation() { RewritePatternSet contiguityPatterns(ctx); populatePimContiguityNormalizationPatterns(contiguityPatterns); - FrozenRewritePatternSet frozenContiguityPatterns(std::move(contiguityPatterns)); - PatternApplicator contiguityApplicator(frozenContiguityPatterns); - contiguityApplicator.applyDefaultCostModel(); - SmallVector contiguityWorklist; - moduleOp.walk([&](Operation* op) { - if (isa(op)) - contiguityWorklist.push_back(op); - }); - - hasFailed = false; - for (Operation* op : contiguityWorklist) { - if (auto copyOp = dyn_cast(op); copyOp && pim::isNormalizedCopyOp(copyOp)) - continue; - if (auto copyOp = dyn_cast(op); copyOp && pim::isNormalizedCopyOp(copyOp)) - continue; - if (auto copyOp = dyn_cast(op); copyOp && pim::isNormalizedCopyOp(copyOp)) - continue; - - if (failed(applyPatternsOnce(op, contiguityApplicator, rewriter))) { - op->emitOpError("failed to normalize PIM copy contiguity"); - hasFailed = true; - } - } - - if (hasFailed) { + GreedyRewriteConfig contiguityConfig; + contiguityConfig.enableFolding(false); + if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) { moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization"); signalPassFailure(); return; @@ -548,13 +465,19 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp (void) walkPimCoreBlockStructurally( coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) { if (auto copyOp = dyn_cast(&op); - copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0))) + copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getSource(), copyOp.getTarget(), + ExpectedPimCopyDirection::DeviceToDevice, + knowledge, failureCount == 0))) ++failureCount; if (auto copyOp = dyn_cast(&op); - copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0))) + copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getHostSource(), copyOp.getDeviceTarget(), + ExpectedPimCopyDirection::HostToDevice, + knowledge, failureCount == 0))) ++failureCount; if (auto copyOp = dyn_cast(&op); - copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0))) + copyOp && failed(verifyPimCopyEndpoints(copyOp, copyOp.getDeviceSource(), copyOp.getHostTarget(), + ExpectedPimCopyDirection::DeviceToHost, + knowledge, failureCount == 0))) ++failureCount; return success(); }); diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp index 7fc2e0e..133e92b 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp @@ -41,7 +41,7 @@ struct HostConstantFoldingPass : PassWrapper patterns; diff --git a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp index a6bb54c..2025cfc 100644 --- a/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp @@ -203,7 +203,7 @@ struct PimMemoryCoalescingPass : PassWrapper buildLocalConcat( ? std::optional(std::move(run)) : std::nullopt; } +static bool canLoopLocalCollection( + const EmitLocalCollectionRun &update, unsigned targetLaneCount) { + if (!update.collection || update.concatenatePayloads + || update.families.size() != 1 + || !(update.lanes == LaneSet::all(targetLaneCount))) + return false; + RequirementFamily &requirement = *update.families.front()->requirement; + return requirement.targetLanes == update.lanes + && !requirement.producerProjection + && (!requirement.producerLocalOffsets + || requirement.producerLocalOffsets->size() == targetLaneCount); +} + +static bool haveSameLocalCollectionContract( + const EmitLocalCollectionRun &lhs, + const EmitLocalCollectionRun &rhs) { + if (lhs.collection != rhs.collection) + return false; + RequirementFamily &left = *lhs.families.front()->requirement; + RequirementFamily &right = *rhs.families.front()->requirement; + if (left.producer->payload != right.producer->payload + || left.publicationFragmentType != right.publicationFragmentType) + return false; + if (lhs.collection->key.kind != FragmentCollectionKind::InsertAssembly) + return true; + const auto &entries = + lhs.collection->key.exchange->program.insertAssembly->entries; + const auto &leftEntry = entries[lhs.collectionPosition]; + const auto &rightEntry = entries[rhs.collectionPosition]; + return leftEntry.sourceTransform == rightEntry.sourceTransform + && leftEntry.sourceType == rightEntry.sourceType; +} + +static void appendLocalUpdates( + BoundaryProgram &boundary, + SmallVectorImpl &updates, + unsigned targetLaneCount) { + for (size_t index = 0; index < updates.size();) { + EmitLocalCollectionRun &first = updates[index]; + if (!canLoopLocalCollection(first, targetLaneCount)) { + boundary.instructions.push_back(std::move(first)); + ++index; + continue; + } + size_t end = index + 1; + while (end < updates.size() + && canLoopLocalCollection(updates[end], targetLaneCount) + && haveSameLocalCollectionContract(first, updates[end])) + ++end; + if (end - index == 1) { + boundary.instructions.push_back(std::move(first)); + ++index; + continue; + } + EmitLocalCollectionLoopRun run; + run.collection = first.collection; + run.lanes = first.lanes; + for (; index < end; ++index) { + run.positions.push_back(updates[index].collectionPosition); + run.families.push_back(updates[index].families.front()); + } + boundary.instructions.push_back(std::move(run)); + } +} + static void appendReceive(BoundaryProgram &boundary, const ScheduledTransferSlice &slice, CollectionTarget target) { @@ -288,8 +353,7 @@ FailureOr buildDeferredBoundaryPlan( if (failed(addCoverage(*local.requirement, local.targetLanes, coverage))) return failure(); } - for (EmitLocalCollectionRun &update : localUpdates) - boundary.instructions.push_back(std::move(update)); + appendLocalUpdates(boundary, localUpdates, exchange->targetLaneCount); for (RequirementFamily &requirement : exchange->requirements) if (!(coverage.lookup(&requirement) == requirement.targetLanes)) return exchange->deferred.emitOpError( diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp index d86a4ef..748374d 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp @@ -20,6 +20,12 @@ struct EmitLocalCollectionRun { LaneSet lanes; bool concatenatePayloads = false; }; +struct EmitLocalCollectionLoopRun { + const FragmentCollectionPlan* collection = nullptr; + llvm::SmallVector positions; + llvm::SmallVector families; + LaneSet lanes; +}; struct EmitReceiveAssemblyRun { const FragmentCollectionPlan* collection = nullptr; llvm::SmallVector slices; @@ -33,7 +39,8 @@ struct ProduceDeferredResult { }; using BoundaryInstruction = - std::variant; struct BoundaryProgram { BoundaryKey key; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp index 220fbc3..731c02f 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp @@ -507,6 +507,137 @@ static FailureOr transformAssemblySource(Value fragment, const DeferredIn llvm_unreachable("unknown deferred assembly source transform"); } +static FailureOr materializeLoopedLocalAssemblySource( + RequirementFamily &requirement, + const DeferredInsertAssemblyEntryTemplate &entry, + Value localOffset, DeferredExchangePlan &exchange, + DeferredEmissionContext &context) { + Value payload = requirement.producer->payload; + auto payloadType = dyn_cast(payload.getType()); + RankedTensorType sourceType = entry.sourceType; + if (entry.sourceTransform + == DeferredAssemblySourceTransform::RemoveLeadingUnitDimension + && payloadType && sourceType + && payloadType.getRank() > sourceType.getRank() + && payloadType.getElementType() == sourceType.getElementType() + && payloadType.getShape().take_back(sourceType.getRank()) + == sourceType.getShape()) { + MixedSliceGeometry geometry; + int64_t rankDifference = payloadType.getRank() - sourceType.getRank(); + geometry.offsets.assign(payloadType.getRank(), + context.rewriter.getIndexAttr(0)); + if (payload.getType() != requirement.publicationFragmentType) + geometry.offsets.front() = localOffset; + geometry.sizes.assign(rankDifference, context.rewriter.getIndexAttr(1)); + for (int64_t dimension : sourceType.getShape()) + geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension)); + geometry.strides.assign(payloadType.getRank(), + context.rewriter.getIndexAttr(1)); + return extractMixedSliceOrIdentity( + context.rewriter, exchange.deferred.getLoc(), payload, sourceType, + geometry); + } + auto fragment = materializeSendPayload( + requirement, localOffset, nullptr, context, + exchange.deferred.getLoc()); + if (failed(fragment)) + return failure(); + return transformAssemblySource(*fragment, entry, exchange, context); +} + +static LogicalResult emitLoopedLocalCollectionUpdate( + const EmitLocalCollectionLoopRun &run, Value lane, unsigned laneCount, + const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) { + if (!run.collection || run.positions.size() < 2 + || run.positions.size() != run.families.size()) + return failure(); + const FragmentCollectionPlan &collection = *run.collection; + DeferredExchangePlan &exchange = *collection.key.exchange; + const DeferredInsertAssemblyEntryTemplate *entry = nullptr; + if (collection.key.kind == FragmentCollectionKind::InsertAssembly) + entry = &exchange.program.insertAssembly + ->entries[run.positions.front()]; + SmallVector offsetRows; + SmallVector positions; + offsetRows.reserve(run.families.size()); + positions.reserve(run.positions.size()); + for (auto [position, family] : llvm::zip_equal(run.positions, + run.families)) { + RequirementFamily &requirement = *family->requirement; + offsetRows.push_back(requirement.producerLocalOffsets + ? *requirement.producerLocalOffsets + : StaticIntSequence::uniform(0, laneCount)); + positions.push_back(position); + } + auto localOffsets = StaticIntGrid::fromRows(offsetRows); + if (failed(localOffsets)) + return failure(); + Operation *anchor = exchange.deferred; + Location loc = anchor->getLoc(); + Value runtimeLane = lane ? lane : context.constants.getIndex(0); + Value current = context.fragmentCollections.lookup(collection.key); + if (!current) + current = createCollectionInitial(collection, context); + auto loop = buildNormalizedScfFor( + context.rewriter, loc, context.constants.getIndex(0), + context.constants.getIndex(run.positions.size()), + context.constants.getIndex(1), ValueRange {current}, + [&](OpBuilder &, Location, Value action, ValueRange iterArgs, + SmallVectorImpl &yielded) -> LogicalResult { + Value position = lookup(positions, action, anchor, context, loc); + Value localOffset = localOffsets->emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc); + RequirementFamily &requirement = *run.families.front()->requirement; + FailureOr source = entry + ? materializeLoopedLocalAssemblySource( + requirement, *entry, localOffset, exchange, context) + : materializeSendPayload( + requirement, localOffset, nullptr, context, loc); + if (failed(source)) + return failure(); + Value next; + if (entry) { + if (source->getType() != entry->sourceType) + return failure(); + next = insertMixedSlice( + context.rewriter, loc, *source, iterArgs.front(), + lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane, + anchor, context, loc)); + } else { + bool grouped = collection.key.kind + == FragmentCollectionKind::GroupedLeaf; + Value specialization = context.constants.getIndex(0); + Value leafPosition = position; + if (grouped) { + Value divisor = context.constants.getIndex(collection.positionCount); + specialization = arith::DivUIOp::create( + context.rewriter, loc, position, divisor); + leafPosition = arith::RemUIOp::create( + context.rewriter, loc, position, divisor); + } + unsigned leafIndex = collection.key.leafIndex; + const DeferredProjectionLeafTemplate &leaf = + exchange.program.leaves[leafIndex]; + auto inserted = insertProjectionFragment( + *source, specialization, leafPosition, + grouped ? specialization : context.constants.getIndex(0), + runtimeLane, iterArgs.front(), leaf, + resultPlan.innerGeometry[leafIndex], exchange, grouped, context); + if (failed(inserted)) + return failure(); + next = *inserted; + } + if (!next) + return failure(); + yielded.push_back(next); + return success(); + }); + if (failed(loop)) + return failure(); + context.fragmentCollections[collection.key] = loop->results.front(); + return success(); +} + static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount, const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) { const FragmentCollectionPlan &collection = *run.collection; @@ -683,6 +814,18 @@ static FailureOr> emitInstructions(ArrayRefdeferred.emitOpError( "failed to update fragment collection from local availability"), failure(); + } else if (auto update = + std::get_if(&instruction)) { + DeferredExchangePlan *exchange = update->collection->key.exchange; + const DeferredResultPlan *resultPlan = findResultPlan(results, exchange); + LogicalResult emitted = resultPlan + ? emitLoopedLocalCollectionUpdate( + *update, lane, laneCount, *resultPlan, context) + : failure(); + if (failed(emitted)) + return exchange->deferred.emitOpError( + "failed to update fragment collection from local assembly run"), + failure(); } else if (auto assembly = std::get_if(&instruction)) { DeferredExchangePlan *exchange = assembly->collection->key.exchange; const DeferredResultPlan *resultPlan = findResultPlan(results, exchange); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp index 9b211e5..4dc6705 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -168,12 +168,16 @@ static LogicalResult materializeResultfulBatchRun( IRMapping mapper; mapper.map(*batch.getLaneArgument(), originalLane); - Value localLane = runLaneCount == 1 - ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0) - : arith::SubIOp::create( - builder, bodyLoc, originalLane, - getOrCreateIndexConstant( - rewriter, batch.getOperation(), first.laneStart)); + Value localLane; + if (runLaneCount == 1) + localLane = getOrCreateIndexConstant(rewriter, batch.getOperation(), 0); + else if (first.laneStart == 0) + localLane = originalLane; + else + localLane = arith::SubIOp::create( + builder, bodyLoc, originalLane, + getOrCreateIndexConstant( + rewriter, batch.getOperation(), first.laneStart)); for (auto [index, weight] : llvm::enumerate(batch.getWeights())) mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); SmallVector inputPlans;