nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Cose belle
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
NiccoloN
2026-07-14 16:57:58 +02:00
parent d1a29ace3c
commit 51fdb830e5
38 changed files with 5365 additions and 914 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp
+5
View File
@@ -400,6 +400,11 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
if (isa<spatial::SpatYieldOp>(op))
continue;
// Cloning a region-bearing operation may leave the rewriter inside that
// region. Every old-block operation is lowered at the core-batch body
// boundary.
rewriter.setInsertionPointToEnd(newBlock);
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
std::optional<StringRef> modeAttr = blueprint.getMode();
if (modeAttr && *modeAttr == "fragment_assembly") {
src/PIM/Conversion/SpatialToPim/Common.cpp
+34 -32
View File
@@ -8,6 +8,8 @@
#include <limits>
#include "Common.hpp"
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
@@ -192,21 +194,23 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
}
static mlir::Value
createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, int64_t stepBytes) {
createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
int64_t stepBytes, Operation *constantAnchor) {
if (stepBytes == 0)
return start;
mlir::Value step = arith::ConstantIndexOp::create(builder, loc, stepBytes);
mlir::Value scaled = arith::MulIOp::create(builder, loc, index, step).getResult();
return arith::AddIOp::create(builder, loc, start, scaled).getResult();
return createOrFoldAffineApply(
builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes,
ValueRange {start, index}, constantAnchor);
}
static mlir::Value createIndexedOffset(OpBuilder& builder,
Location loc,
mlir::Value indexArg,
ArrayRef<int64_t> values) {
ArrayRef<int64_t> values,
Operation *constantAnchor) {
assert(!values.empty() && "expected lane-indexed values");
if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); }))
return arith::ConstantIndexOp::create(builder, loc, values.front());
return getOrCreateIndexConstant(builder, constantAnchor, values.front());
if (values.size() >= 2) {
int64_t step = values[1] - values[0];
@@ -214,21 +218,18 @@ static mlir::Value createIndexedOffset(OpBuilder& builder,
return values[index] == values.front() + static_cast<int64_t>(index) * step;
});
if (arithmetic) {
mlir::Value base = arith::ConstantIndexOp::create(builder, loc, values.front());
mlir::Value stepValue = arith::ConstantIndexOp::create(builder, loc, step);
mlir::Value scaledIndex = arith::MulIOp::create(builder, loc, indexArg, stepValue).getResult();
return arith::AddIOp::create(builder, loc, base, scaledIndex).getResult();
return createOrFoldAffineApply(
builder, loc, builder.getAffineDimExpr(0) * step + values.front(),
ValueRange {indexArg}, constantAnchor);
}
}
mlir::Value selected = arith::ConstantIndexOp::create(builder, loc, values.front());
for (auto [lane, value] : llvm::enumerate(values.drop_front())) {
mlir::Value indexValue = arith::ConstantIndexOp::create(builder, loc, static_cast<int64_t>(lane + 1));
mlir::Value cmp = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq, indexArg, indexValue);
mlir::Value candidate = arith::ConstantIndexOp::create(builder, loc, value);
selected = arith::SelectOp::create(builder, loc, cmp, candidate, selected);
}
return selected;
RankedTensorType tableType = RankedTensorType::get(
{static_cast<int64_t>(values.size())}, builder.getI64Type());
DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg});
return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult();
}
struct FragmentAssemblyCopyRunFamily {
@@ -433,11 +434,11 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
mlir::Value hostStart;
mlir::Value sourceStart;
if (laneArg) {
hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane);
sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane);
hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor);
sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor);
} else {
hostStart = arith::ConstantIndexOp::create(builder, loc, run.hostStartBytesByLane.front());
sourceStart = arith::ConstantIndexOp::create(builder, loc, run.sourceStartBytesByLane.front());
hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front());
sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front());
}
if (hostRunStartDelta)
@@ -459,9 +460,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
.getOutput();
}
mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, run.count);
mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count);
mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
builder,
loc,
@@ -474,9 +475,10 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
mlir::Value flatIndex,
ValueRange iterArgs,
SmallVectorImpl<mlir::Value>& yielded) {
mlir::Value hostOffset = createSteppedOffset(loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes);
mlir::Value hostOffset = createSteppedOffset(
loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor);
mlir::Value sourceOffset =
createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes);
createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor);
mlir::Value copied =
pim::PimMemCopyDevToHostOp::create(loopBuilder,
bodyLoc,
@@ -506,9 +508,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
return emitFragmentAssemblyCopyRun(
builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset);
mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, family.sourceRunStartDeltas.size());
mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size());
mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
builder,
loc,
@@ -522,9 +524,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
ValueRange iterArgs,
SmallVectorImpl<mlir::Value>& yielded) {
mlir::Value sourceRunStartDelta =
createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas);
createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor);
mlir::Value hostRunStartDelta =
createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas);
createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor);
FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
bodyLoc,
family.prototype,
src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp
+83 -1
View File
@@ -10,7 +10,9 @@
#include "Conversion/ONNXToSpatial/Common/Common.hpp"
#include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -180,16 +182,79 @@ static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute com
return success();
}
static bool isHostMaterializableHelperOp(Operation* op) {
if (isa<spatial::SpatYieldOp>(op))
return true;
if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
return true;
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
std::optional<StringRef> mode = blueprint.getMode();
return mode && *mode == "fragment_assembly";
}
return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
}
static FailureOr<DenseMap<Value, Attribute>>
analyzeHostMaterializableHelper(spatial::SpatScheduledCompute computeOp) {
DenseMap<Value, Attribute> folded;
for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights())) {
auto argument = computeOp.getWeightArgument(weightIndex);
if (!argument)
return failure();
Attribute constant;
if (matchPattern(weight, m_Constant(&constant)))
folded[*argument] = constant;
}
Block& block = computeOp.getBody().front();
for (Operation& op : block) {
if (!isHostMaterializableHelperOp(&op))
return failure();
if (isa<spatial::SpatYieldOp, spatial::SpatBlueprintOp>(op)
|| (isShapingOnlyOp(&op) && !isPureIndexComputationOp(&op)))
continue;
if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()) {
for (Value result : op.getResults()) {
Attribute constant;
if (!matchPattern(result, m_Constant(&constant)))
return failure();
folded[result] = constant;
}
continue;
}
if (!isPureIndexComputationOp(&op) || op.getNumRegions() != 0)
return failure();
SmallVector<Attribute> operands;
for (Value operand : op.getOperands()) {
auto it = folded.find(operand);
if (it == folded.end())
return failure();
operands.push_back(it->second);
}
SmallVector<OpFoldResult> results;
if (failed(op.fold(operands, results))
|| results.size() != op.getNumResults())
return failure();
for (auto [result, foldResult] : llvm::zip(op.getResults(), results)) {
auto attribute = dyn_cast<Attribute>(foldResult);
if (!attribute)
return failure();
folded[result] = attribute;
}
}
return folded;
}
static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp,
IRRewriter& rewriter,
OperationFolder& constantFolder) {
if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
return false;
if (computeOp.getResult(0).use_empty())
return false;
if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
return isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
}))
return false;
Block& block = computeOp.getBody().front();
if (block.getNumArguments() != computeOp.getWeights().size())
return false;
@@ -197,6 +262,9 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
if (!yieldOp || yieldOp.getNumOperands() != 1)
return false;
auto folded = analyzeHostMaterializableHelper(computeOp);
if (failed(folded))
return false;
rewriter.setInsertionPoint(computeOp);
IRMapping mapping;
@@ -218,6 +286,20 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
}
}
if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()
|| isPureIndexComputationOp(&op)) {
for (Value result : op.getResults()) {
auto it = folded->find(result);
if (it == folded->end())
return false;
mapping.map(
result,
getOrCreateConstant(constantFolder, computeOp, it->second,
result.getType()));
}
continue;
}
cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
Operation* clonedOp = rewriter.clone(op, mapping);
for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))