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

Merge branch 'TestRottoConDeadLock' of chef.heaplab.deib.polimi.it:nnicolosi/Raptor into TestRottoConDeadLock
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source
This commit is contained in:
ilgeco
2026-07-20 18:06:55 +02:00
parent 961fd613bd 2d17b50ca8
commit a893d23a74
28 changed files with 1378 additions and 1241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Conversion/SpatialToPim/Common.cpp
+23 -1
View File
@@ -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<tensor::CastOp,
tensor::CollapseShapeOp,
tensor::ExpandShapeOp,
tensor::ExtractSliceOp,
tensor::ReshapeOp>(value.getDefiningOp());
}
static bool isLoopCarriedOutput(mlir::Value operand, Operation* operation) {
auto argument = dyn_cast<BlockArgument>(operand);
if (!argument || argument.getArgNumber() == 0 || operation->getBlock() != argument.getOwner())
return false;
auto loop = dyn_cast_or_null<scf::ForOp>(argument.getOwner()->getParentOp());
return loop && cast<scf::YieldOp>(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<mlir::Value> 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<BlockArgument>(operand) || isLoopCarriedOutput(operand, operation))
&& !operand.getDefiningOp<arith::ConstantOp>()
&& !isTensorView(operand)
&& !hasLaterUserInBlock(operand, operation);
});
auto bestOperand = validOperands.begin();
src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp
+8 -3
View File
@@ -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<spatial::SpatConcatOp,
target.addLegalOp<linalg::MapOp,
linalg::YieldOp,
spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,
spatial::SpatExtractRowsOp>();
@@ -186,7 +189,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
memref::MemRefDialect,
scf::SCFDialect,
BuiltinDialect>();
coreBodyTarget.addLegalOp<spatial::SpatConcatOp,
coreBodyTarget.addLegalOp<linalg::MapOp,
linalg::YieldOp,
spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,
spatial::SpatExtractRowsOp>();
@@ -234,7 +239,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
memref::MemRefDialect,
scf::SCFDialect,
BuiltinDialect>();
communicationTarget.addLegalOp<ModuleOp>();
communicationTarget.addLegalOp<ModuleOp, linalg::MapOp, linalg::YieldOp>();
communicationTarget.addIllegalOp<spatial::SpatConcatOp,
spatial::SpatChannelReceiveOp,
spatial::SpatChannelSendOp,