fixes

slightly faster codegen

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -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();