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

finish helper refactoring
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source 
use uniqued constant helpers everywhere
materialize transposed constants directly
This commit is contained in:
NiccoloN
2026-05-29 17:05:45 +02:00
parent 819d8af0f7
commit 8bb0babf1b
32 changed files with 300 additions and 467 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Pass/PimCodegen/HostConstantFolding/Patterns/Subview.cpp
+11 -10
View File
@@ -2,6 +2,7 @@
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "../Common.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -41,7 +42,7 @@ static OpFoldResult addConstantOffset(OpFoldResult baseOffset, int64_t extraOffs
}
auto value = cast<Value>(baseOffset);
auto cst = arith::ConstantIndexOp::create(rewriter, value.getLoc(), extraOffset);
auto cst = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), extraOffset);
return arith::AddIOp::create(rewriter, value.getLoc(), value, cst).getResult();
}
@@ -75,7 +76,7 @@ static SmallVector<Value> delinearizeIndexValue(Value linearIndex,
Value remaining = linearIndex;
for (auto [_dim, stride] : llvm::enumerate(strides)) {
auto cStride = arith::ConstantIndexOp::create(rewriter, linearIndex.getLoc(), stride);
auto cStride = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), stride);
Value index = arith::DivUIOp::create(rewriter, linearIndex.getLoc(), remaining, cStride);
indices.push_back(index);
remaining = arith::RemUIOp::create(rewriter, linearIndex.getLoc(), remaining, cStride);
@@ -90,7 +91,7 @@ static OpFoldResult addDynamicOffset(OpFoldResult baseOffset, Value extraOffset,
if (integerAttr.getInt() == 0)
return extraOffset;
auto cst = arith::ConstantIndexOp::create(rewriter, extraOffset.getLoc(), integerAttr.getInt());
auto cst = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), integerAttr.getInt());
return arith::AddIOp::create(rewriter, extraOffset.getLoc(), cst, extraOffset).getResult();
}
@@ -195,9 +196,9 @@ static LogicalResult rewriteSubviewCopyLikeOp(CopyOp copyOp,
if (allowLoopRewrite && numSlices > 1 && srcOffset == 0 && dstOffset == 0
&& sourceType.getRank() == static_cast<int64_t>(copyShape.size())
&& dstType.getRank() == static_cast<int64_t>(copyShape.size())) {
auto c0 = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), 0);
auto cUpper = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), numSlices);
auto cStep = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), 1);
auto c0 = getOrCreateIndexConstant(rewriter, copyOp, 0);
auto cUpper = getOrCreateIndexConstant(rewriter, copyOp, numSlices);
auto cStep = getOrCreateIndexConstant(rewriter, copyOp, 1);
auto loop = scf::ForOp::create(rewriter, copyOp.getLoc(), c0, cUpper, cStep, ValueRange {});
rewriter.setInsertionPointToStart(loop.getBody());
@@ -284,8 +285,8 @@ struct RewriteHostSubviewLoadPattern final : OpRewritePattern<pim::PimMemCopyHos
rewriter,
[&](
MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
Value dstOffsetValue = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), dstByteOffset);
Value srcOffsetValue = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), srcByteOffset);
Value dstOffsetValue = getOrCreateIndexConstant(rewriter, copyOp, dstByteOffset);
Value srcOffsetValue = getOrCreateIndexConstant(rewriter, copyOp, srcByteOffset);
pim::PimMemCopyHostToDevOp::create(rewriter,
copyOp.getLoc(),
resultType,
@@ -324,8 +325,8 @@ struct RewriteHostSubviewStorePattern final : OpRewritePattern<pim::PimMemCopyDe
rewriter,
[&](
MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
Value dstOffset = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), dstByteOffset);
Value srcOffset = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), srcByteOffset);
Value dstOffset = getOrCreateIndexConstant(rewriter, copyOp, dstByteOffset);
Value srcOffset = getOrCreateIndexConstant(rewriter, copyOp, srcByteOffset);
pim::PimMemCopyDevToHostOp::create(rewriter,
copyOp.getLoc(),
resultType,
src/PIM/Pass/PimCodegen/MaterializeHostConstantsPass.cpp
+4 -13
View File
@@ -13,6 +13,7 @@
#include <type_traits>
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -22,16 +23,6 @@ namespace onnx_mlir {
namespace {
static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
if (isa<pim::PimMemCopyHostToDevOp>(op))
return operandIndex == 3;
if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
return operandIndex == 1;
if (isa<pim::PimMemCopyDevToHostOp>(op))
return operandIndex == 2;
return false;
}
template <typename CoreOpTy>
static void materializeHostConstantsInCore(CoreOpTy coreOp,
IRRewriter& rewriter,
@@ -51,7 +42,7 @@ static void materializeHostConstantsInCore(CoreOpTy coreOp,
for (OpOperand& operand : op->getOpOperands()) {
Value originalValue = operand.get();
if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostOperand(op, operand.getOperandNumber()))
if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostMemCopyOperand(op, operand.getOperandNumber()))
continue;
auto resolvedAddress = resolveContiguousAddress(originalValue);
@@ -113,8 +104,8 @@ static void materializeHostConstantsInCore(CoreOpTy coreOp,
rewriter,
op->getLoc(),
originalType,
getOrCreateHostIndexConstant(constantFolder, op, 0),
getOrCreateHostIndexConstant(constantFolder, op, static_cast<int64_t>(resolvedAddress->byteOffset) ),
getOrCreateIndexConstant(constantFolder, op, 0),
getOrCreateIndexConstant(constantFolder, op, static_cast<int64_t>(resolvedAddress->byteOffset) ),
deviceDst,
getGlobalOp.getResult(),
rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
src/PIM/Pass/PimCodegen/VerificationPass.cpp
+2 -11
View File
@@ -6,6 +6,7 @@
#include "llvm/ADT/STLExtras.h"
#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
@@ -119,16 +120,6 @@ static bool isConstantGlobalView(Value value) {
}
}
static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
if (isa<pim::PimMemCopyHostToDevOp>(op))
return operandIndex == 3;
if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
return operandIndex == 1;
if (isa<pim::PimMemCopyDevToHostOp>(op))
return operandIndex == 2;
return false;
}
static bool isCoreWeightBlockArgument(Value value) {
auto blockArgument = dyn_cast<BlockArgument>(value);
if (!blockArgument)
@@ -361,7 +352,7 @@ private:
continue;
}
if (isExplicitHostOperand(&op, operandIndex)) {
if (isExplicitHostMemCopyOperand(&op, operandIndex)) {
if (!isCodegenAddressableValue(operand, knowledge)) {
diagnostics.report(&op, [&](Operation* illegalOp) {
illegalOp->emitOpError() << "host operand #" << operandIndex