fixes

slightly faster codegen
2026-07-20 16:01:57 +02:00
parent ab54243fda
commit dbb66be93e
27 changed files with 1358 additions and 1241 deletions
@@ -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();

@@ -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,