nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 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/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,