Fix raising with index cast

include/polygeist/BarrierUtils.h

@@ -74,6 +74,23 @@ static T allocateTemporaryBuffer(mlir::OpBuilder &rewriter, mlir::Value value,
   return rewriter.create<T>(value.getLoc(), type, iterationCounts);
 }
 
+template <>
+mlir::LLVM::AllocaOp allocateTemporaryBuffer<mlir::LLVM::AllocaOp>(
+    mlir::OpBuilder &rewriter, mlir::Value value,
+    mlir::ValueRange iterationCounts, bool alloca, mlir::DataLayout *DLI) {
+  using namespace mlir;
+  auto val = value.getDefiningOp<LLVM::AllocaOp>();
+  auto sz = val.getArraySize();
+  assert(DLI);
+  for (auto iter : iterationCounts) {
+    sz =
+        rewriter.create<arith::MulIOp>(value.getLoc(), sz,
+                                       rewriter.create<arith::IndexCastOp>(
+                                           value.getLoc(), sz.getType(), iter));
+  }
+  return rewriter.create<LLVM::AllocaOp>(value.getLoc(), val.getType(), sz);
+}
+
 template <>
 mlir::LLVM::CallOp allocateTemporaryBuffer<mlir::LLVM::CallOp>(
     mlir::OpBuilder &rewriter, mlir::Value value,

include/polygeist/Passes/Passes.h

@@ -25,7 +25,8 @@ std::unique_ptr<Pass> createConvertPolygeistToLLVMPass();
 } // namespace mlir
 
 void fully2ComposeAffineMapAndOperands(
-    mlir::AffineMap *map, llvm::SmallVectorImpl<mlir::Value> *operands);
+    mlir::OpBuilder &, mlir::AffineMap *map,
+    llvm::SmallVectorImpl<mlir::Value> *operands);
 bool isValidIndex(mlir::Value val);
 
 namespace mlir {

include/polygeist/Passes/Passes.td

@@ -8,7 +8,7 @@ def AffineCFG : Pass<"affine-cfg"> {
   let constructor = "mlir::polygeist::replaceAffineCFGPass()";
 }
 
-def Mem2Reg : Pass<"mem2reg", "FuncOp"> {
+def Mem2Reg : Pass<"mem2reg"> {
   let summary = "Replace scf.if and similar with affine.if";
   let constructor = "mlir::polygeist::createMem2RegPass()";
 }
@@ -25,7 +25,7 @@ def AffineReduction : Pass<"detect-reduction"> {
   let constructor = "mlir::polygeist::detectReductionPass()";
 }
 
-def SCFCPUify : Pass<"cpuify", "FuncOp"> {
+def SCFCPUify : Pass<"cpuify"> {
   let summary = "remove scf.barrier";
   let constructor = "mlir::polygeist::createCPUifyPass()";
   let dependentDialects =
@@ -35,7 +35,7 @@ def SCFCPUify : Pass<"cpuify", "FuncOp"> {
   ];
 }
 
-def SCFBarrierRemovalContinuation : Pass<"barrier-removal-continuation", "FuncOp"> {
+def SCFBarrierRemovalContinuation : InterfacePass<"barrier-removal-continuation", "FunctionOpInterface"> {
   let summary = "Remove scf.barrier using continuations";
   let constructor = "mlir::polygeist::createBarrierRemovalContinuation()";
   let dependentDialects = ["memref::MemRefDialect", "func::FuncDialect"];

include/polygeist/Passes/Utils.h

@@ -0,0 +1,109 @@
+#pragma once
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/IntegerSet.h"
+
+static inline mlir::scf::IfOp
+cloneWithoutResults(mlir::scf::IfOp op, mlir::OpBuilder &rewriter,
+                    mlir::BlockAndValueMapping mapping = {},
+                    mlir::TypeRange types = {}) {
+  using namespace mlir;
+  return rewriter.create<scf::IfOp>(
+      op.getLoc(), types, mapping.lookupOrDefault(op.getCondition()), true);
+}
+static inline mlir::AffineIfOp
+cloneWithoutResults(mlir::AffineIfOp op, mlir::OpBuilder &rewriter,
+                    mlir::BlockAndValueMapping mapping = {},
+                    mlir::TypeRange types = {}) {
+  using namespace mlir;
+  SmallVector<mlir::Value> lower;
+  for (auto o : op.getOperands())
+    lower.push_back(mapping.lookupOrDefault(o));
+  return rewriter.create<AffineIfOp>(op.getLoc(), types, op.getIntegerSet(),
+                                     lower, true);
+}
+
+static inline mlir::scf::ForOp
+cloneWithoutResults(mlir::scf::ForOp op, mlir::PatternRewriter &rewriter,
+                    mlir::BlockAndValueMapping mapping = {}) {
+  using namespace mlir;
+  return rewriter.create<scf::ForOp>(
+      op.getLoc(), mapping.lookupOrDefault(op.getLowerBound()),
+      mapping.lookupOrDefault(op.getUpperBound()),
+      mapping.lookupOrDefault(op.getStep()));
+}
+static inline mlir::AffineForOp
+cloneWithoutResults(mlir::AffineForOp op, mlir::PatternRewriter &rewriter,
+                    mlir::BlockAndValueMapping mapping = {}) {
+  using namespace mlir;
+  SmallVector<Value> lower;
+  for (auto o : op.getLowerBoundOperands())
+    lower.push_back(mapping.lookupOrDefault(o));
+  SmallVector<Value> upper;
+  for (auto o : op.getUpperBoundOperands())
+    upper.push_back(mapping.lookupOrDefault(o));
+  return rewriter.create<AffineForOp>(op.getLoc(), lower, op.getLowerBoundMap(),
+                                      upper, op.getUpperBoundMap(),
+                                      op.getStep());
+}
+
+static inline mlir::Block *getThenBlock(mlir::scf::IfOp op) {
+  return op.thenBlock();
+}
+static inline mlir::Block *getThenBlock(mlir::AffineIfOp op) {
+  return op.getThenBlock();
+}
+static inline mlir::Block *getElseBlock(mlir::scf::IfOp op) {
+  return op.elseBlock();
+}
+static inline mlir::Block *getElseBlock(mlir::AffineIfOp op) {
+  return op.getElseBlock();
+}
+
+static inline mlir::Region &getThenRegion(mlir::scf::IfOp op) {
+  return op.getThenRegion();
+}
+static inline mlir::Region &getThenRegion(mlir::AffineIfOp op) {
+  return op.thenRegion();
+}
+static inline mlir::Region &getElseRegion(mlir::scf::IfOp op) {
+  return op.getElseRegion();
+}
+static inline mlir::Region &getElseRegion(mlir::AffineIfOp op) {
+  return op.elseRegion();
+}
+
+static inline mlir::scf::YieldOp getThenYield(mlir::scf::IfOp op) {
+  return op.thenYield();
+}
+static inline mlir::AffineYieldOp getThenYield(mlir::AffineIfOp op) {
+  return llvm::cast<mlir::AffineYieldOp>(op.getThenBlock()->getTerminator());
+}
+static inline mlir::scf::YieldOp getElseYield(mlir::scf::IfOp op) {
+  return op.elseYield();
+}
+static inline mlir::AffineYieldOp getElseYield(mlir::AffineIfOp op) {
+  return llvm::cast<mlir::AffineYieldOp>(op.getElseBlock()->getTerminator());
+}
+
+static inline bool inBound(mlir::scf::IfOp op, mlir::Value v) {
+  return op.getCondition() == v;
+}
+static inline bool inBound(mlir::AffineIfOp op, mlir::Value v) {
+  return llvm::any_of(op.getOperands(), [&](mlir::Value e) { return e == v; });
+}
+static inline bool inBound(mlir::scf::ForOp op, mlir::Value v) {
+  return op.getUpperBound() == v;
+}
+static inline bool inBound(mlir::AffineForOp op, mlir::Value v) {
+  return llvm::any_of(op.getUpperBoundOperands(),
+                      [&](mlir::Value e) { return e == v; });
+}
+static inline bool hasElse(mlir::scf::IfOp op) {
+  return op.getElseRegion().getBlocks().size() > 0;
+}
+static inline bool hasElse(mlir::AffineIfOp op) {
+  return op.elseRegion().getBlocks().size() > 0;
+}

include/polygeist/PolygeistOps.td

@@ -24,6 +24,7 @@ def BarrierOp
   let arguments = (ins Variadic<Index>:$indices);
   let summary = "barrier for parallel loops";
   let description = [{}];
+  let hasCanonicalizer = true;
 }
 
 //===----------------------------------------------------------------------===//