Alternative lowering for memrefs (#276)

* WIP: alternative lowering for memrefs

Introduce an alternative lowering scheme for memref types. A n-D memref
is always lowered to a pointer to an (n-1)-D nested LLVM array type, a
0-D and a 1-D memref is lowered to a bare pointer. This applies both at
function boundaries and within functions, unlike the MLIR's bare pointer
calling convention. This lowering is only possible for memrefs with
static shapes except for the outermost dimension, which is consistent
with the kinds of nested array structure C allows and therefore suitable
for a C frontend. However, it only supports a subset of memref dialect
operations that are possible to implement using this data structure. For
example, view, subview and transpose cannot be implemented as they
result in non-row-major layouts unexpressible with C nested arrays.

This simplifies the lowering and the generated IR for all supported
memref operations and removes the need for the ABI duplication at the
function boundary: memrefs are lowered to the same LLVM IR as equivalent
nested arrays.

* Enable c memref

* Fix lowering c abi

* Fix compile issues

* Fix format

* Add unit tests and fix the bug discovered in the process

nD dynamic allocation was not taking into account the static part of the
shape leading to a smaller allocation than required.

* Update tests

Golden tests are simply rewritten with the version. The main difference
is the pervasiveness of `memref` in the function signatures and the
corresponding memref-to-pointer conversions (or lack thereof) around
function calls.

* Create a pointer2memref op if we cast between different types

* Do not use argv special case

* Fix bug with multiple definitions of malloc

* subindexlowering: bitcast after gepop

* Update test

* Bugfixes

* Cleanup

* Fix LLVM bump

* Fix tests

* Generalize tests

Co-authored-by: William S. Moses <[email protected]>
Co-authored-by: Ivan Radanov Ivanov <[email protected]>

Author: 3 people

include/polygeist/BarrierUtils.h

@@ -17,7 +17,7 @@
 #include "mlir/IR/Block.h"
 #include "polygeist/Ops.h"
 #include "llvm/ADT/SetVector.h"
-#include <mlir/Dialect/Arithmetic/IR/Arithmetic.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
 
 std::pair<mlir::Block *, mlir::Block::iterator>
 findInsertionPointAfterLoopOperands(mlir::scf::ParallelOp op);
@@ -41,14 +41,15 @@ emitIterationCounts(mlir::OpBuilder &rewriter, mlir::scf::ParallelOp op) {
   return iterationCounts;
 }
 
-mlir::LLVM::LLVMFuncOp GetOrCreateMallocFunction(mlir::ModuleOp module);
+mlir::Value callMalloc(mlir::OpBuilder &builder, mlir::ModuleOp module,
+                       mlir::Location loc, mlir::Value arg);
 mlir::LLVM::LLVMFuncOp GetOrCreateFreeFunction(mlir::ModuleOp module);
 
 template <typename T>
-static T allocateTemporaryBuffer(mlir::OpBuilder &rewriter, mlir::Value value,
-                                 mlir::ValueRange iterationCounts,
-                                 bool alloca = true,
-                                 mlir::DataLayout *DLI = nullptr) {
+static mlir::Value
+allocateTemporaryBuffer(mlir::OpBuilder &rewriter, mlir::Value value,
+                        mlir::ValueRange iterationCounts, bool alloca = true,
+                        mlir::DataLayout *DLI = nullptr) {
   using namespace mlir;
   SmallVector<int64_t> bufferSize(iterationCounts.size(),
                                   ShapedType::kDynamicSize);
@@ -75,7 +76,7 @@ static T allocateTemporaryBuffer(mlir::OpBuilder &rewriter, mlir::Value value,
 }
 
 template <>
-mlir::LLVM::AllocaOp allocateTemporaryBuffer<mlir::LLVM::AllocaOp>(
+mlir::Value allocateTemporaryBuffer<mlir::LLVM::AllocaOp>(
     mlir::OpBuilder &rewriter, mlir::Value value,
     mlir::ValueRange iterationCounts, bool alloca, mlir::DataLayout *DLI) {
   using namespace mlir;
@@ -92,7 +93,7 @@ mlir::LLVM::AllocaOp allocateTemporaryBuffer<mlir::LLVM::AllocaOp>(
 }
 
 template <>
-mlir::LLVM::CallOp allocateTemporaryBuffer<mlir::LLVM::CallOp>(
+mlir::Value allocateTemporaryBuffer<mlir::LLVM::CallOp>(
     mlir::OpBuilder &rewriter, mlir::Value value,
     mlir::ValueRange iterationCounts, bool alloca, mlir::DataLayout *DLI) {
   using namespace mlir;
@@ -113,7 +114,7 @@ mlir::LLVM::CallOp allocateTemporaryBuffer<mlir::LLVM::CallOp>(
                                            value.getLoc(), sz.getType(), iter));
   }
   auto m = val->getParentOfType<ModuleOp>();
-  auto allocfn = GetOrCreateMallocFunction(m);
-  return rewriter.create<LLVM::CallOp>(value.getLoc(), allocfn, sz);
+  return callMalloc(rewriter, m, value.getLoc(), sz);
 }
+
 #endif // MLIR_LIB_DIALECT_SCF_TRANSFORMS_BARRIERUTILS_H_

include/polygeist/Passes/Passes.h

@@ -24,7 +24,8 @@ std::unique_ptr<Pass> detectReductionPass();
 std::unique_ptr<Pass> createRemoveTrivialUsePass();
 std::unique_ptr<Pass> createParallelLowerPass();
 std::unique_ptr<Pass>
-createConvertPolygeistToLLVMPass(const LowerToLLVMOptions &options);
+createConvertPolygeistToLLVMPass(const LowerToLLVMOptions &options,
+                                 bool useCStyleMemRef);
 std::unique_ptr<Pass> createConvertPolygeistToLLVMPass();
 std::unique_ptr<Pass> createForBreakToWhilePass();
 
@@ -43,7 +44,7 @@ template <typename ConcreteDialect>
 void registerDialect(DialectRegistry &registry);
 
 namespace arith {
-class ArithmeticDialect;
+class ArithDialect;
 } // end namespace arith
 
 namespace scf {

include/polygeist/Passes/Passes.td

@@ -68,7 +68,7 @@ def SCFCanonicalizeFor : Pass<"canonicalize-scf-for"> {
 def ForBreakToWhile : Pass<"for-break-to-while"> {
   let summary = "Rewrite scf.for(scf.if) to scf.while";
   let constructor = "mlir::polygeist::createForBreakToWhilePass()";
-  let dependentDialects = ["arith::ArithmeticDialect"];
+  let dependentDialects = ["arith::ArithDialect"];
 }
 
 def ParallelLICM : Pass<"parallel-licm"> {
@@ -130,7 +130,11 @@ def ConvertPolygeistToLLVM : Pass<"convert-polygeist-to-llvm", "mlir::ModuleOp">
     Option<"dataLayout", "data-layout", "std::string",
            /*default=*/"\"\"",
            "String description (LLVM format) of the data layout that is "
-           "expected on the produced module">
+           "expected on the produced module">,
+    Option<"useCStyleMemRef", "use-c-style-memref", "bool",
+           /*default=*/"true",
+           "Use C-style nested-array lowering of memref instead of "
+           "the default MLIR descriptor structure">
   ];
 }
 

lib/polygeist/Ops.cpp

@@ -7,7 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "polygeist/Ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/AffineExpr.h"
@@ -20,7 +20,7 @@
 #include "polygeist/PolygeistOps.cpp.inc"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -384,7 +384,7 @@ static bool mayAlias(Value v, Value v2) {
 
   if (auto glob = v.getDefiningOp<memref::GetGlobalOp>()) {
     if (auto Aglob = v2.getDefiningOp<memref::GetGlobalOp>()) {
-      return glob.name() == Aglob.name();
+      return glob.getName() == Aglob.getName();
     }
   }
 
@@ -644,7 +644,7 @@ struct SimplifySubIndexUsers : public OpRewritePattern<SubIndexOp> {
                                                        subindex.getSource());
       } else if (auto loadOp = dyn_cast<memref::LoadOp>(use.getOwner())) {
         if (loadOp.getMemref() == subindex) {
-          SmallVector<Value, 4> indices = loadOp.indices();
+          SmallVector<Value, 4> indices = loadOp.getIndices();
           if (subindex.getType().cast<MemRefType>().getShape().size() ==
               subindex.getSource()
                   .getType()
@@ -676,7 +676,7 @@ struct SimplifySubIndexUsers : public OpRewritePattern<SubIndexOp> {
         }
       } else if (auto storeOp = dyn_cast<memref::StoreOp>(use.getOwner())) {
         if (storeOp.getMemref() == subindex) {
-          SmallVector<Value, 4> indices = storeOp.indices();
+          SmallVector<Value, 4> indices = storeOp.getIndices();
           if (subindex.getType().cast<MemRefType>().getShape().size() ==
               subindex.getSource()
                   .getType()
@@ -707,7 +707,7 @@ struct SimplifySubIndexUsers : public OpRewritePattern<SubIndexOp> {
         }
       } else if (auto storeOp = dyn_cast<memref::AtomicRMWOp>(use.getOwner())) {
         if (storeOp.getMemref() == subindex) {
-          SmallVector<Value, 4> indices = storeOp.indices();
+          SmallVector<Value, 4> indices = storeOp.getIndices();
           if (subindex.getType().cast<MemRefType>().getShape().size() ==
               subindex.getSource()
                   .getType()
@@ -733,7 +733,7 @@ struct SimplifySubIndexUsers : public OpRewritePattern<SubIndexOp> {
                      .getShape()
                      .size() == indices.size());
           rewriter.replaceOpWithNewOp<memref::AtomicRMWOp>(
-              storeOp, storeOp.getType(), storeOp.kind(), storeOp.getValue(),
+              storeOp, storeOp.getType(), storeOp.getKind(), storeOp.getValue(),
               subindex.getSource(), indices);
           changed = true;
         }
@@ -843,7 +843,7 @@ struct SimplifySubViewUsers : public OpRewritePattern<memref::SubViewOp> {
                                                        subindex.getSource());
       } else if (auto loadOp = dyn_cast<memref::LoadOp>(use.getOwner())) {
         if (loadOp.getMemref() == subindex) {
-          SmallVector<Value, 4> indices = loadOp.indices();
+          SmallVector<Value, 4> indices = loadOp.getIndices();
           if (subindex.getType().cast<MemRefType>().getShape().size() ==
               subindex.getSource()
                   .getType()
@@ -878,7 +878,7 @@ struct SimplifySubViewUsers : public OpRewritePattern<memref::SubViewOp> {
         }
       } else if (auto storeOp = dyn_cast<memref::StoreOp>(use.getOwner())) {
         if (storeOp.getMemref() == subindex) {
-          SmallVector<Value, 4> indices = storeOp.indices();
+          SmallVector<Value, 4> indices = storeOp.getIndices();
           if (subindex.getType().cast<MemRefType>().getShape().size() ==
               subindex.getSource()
                   .getType()
@@ -1539,7 +1539,7 @@ class MetaPointer2Memref final : public OpRewritePattern<Op> {
 template <>
 Value MetaPointer2Memref<memref::LoadOp>::computeIndex(
     memref::LoadOp op, size_t i, PatternRewriter &rewriter) const {
-  return op.indices()[i];
+  return op.getIndices()[i];
 }
 
 template <>
@@ -1551,7 +1551,7 @@ void MetaPointer2Memref<memref::LoadOp>::rewrite(
 template <>
 Value MetaPointer2Memref<memref::StoreOp>::computeIndex(
     memref::StoreOp op, size_t i, PatternRewriter &rewriter) const {
-  return op.indices()[i];
+  return op.getIndices()[i];
 }
 
 template <>
@@ -2551,15 +2551,15 @@ struct ConstantRankReduction : public OpRewritePattern<memref::AllocaOp> {
     for (auto u : op->getResult(0).getUsers()) {
       if (auto load = dyn_cast<memref::LoadOp>(u)) {
         if (!set) {
-          for (auto i : load.indices()) {
+          for (auto i : load.getIndices()) {
             IntegerAttr constValue;
             if (!matchPattern(i, m_Constant(&constValue)))
               return failure();
             v.push_back(constValue.getValue().getZExtValue());
           }
           set = true;
         } else {
-          for (auto pair : llvm::zip(load.indices(), v)) {
+          for (auto pair : llvm::zip(load.getIndices(), v)) {
             IntegerAttr constValue;
             if (!matchPattern(std::get<0>(pair), m_Constant(&constValue)))
               return failure();
@@ -2626,7 +2626,7 @@ struct ConstantRankReduction : public OpRewritePattern<memref::AllocaOp> {
       }
       if (auto store = dyn_cast<memref::StoreOp>(u)) {
         Value cond = nullptr;
-        for (auto pair : llvm::zip(store.indices(), v)) {
+        for (auto pair : llvm::zip(store.getIndices(), v)) {
           auto val = rewriter.create<arith::CmpIOp>(
               store.getLoc(), CmpIPredicate::eq, std::get<0>(pair),
               rewriter.create<arith::ConstantIndexOp>(store.getLoc(),

lib/polygeist/Passes/AffineCFG.cpp

@@ -15,7 +15,7 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Debug.h"
 #include <deque>
-#include <mlir/Dialect/Arithmetic/IR/Arithmetic.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
 
 #define DEBUG_TYPE "affine-cfg"
 

lib/polygeist/Passes/BarrierRemovalContinuation.cpp

@@ -14,7 +14,7 @@
 #include "PassDetails.h"
 
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"

lib/polygeist/Passes/CMakeLists.txt

@@ -24,7 +24,7 @@ add_mlir_dialect_library(MLIRPolygeistTransforms
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
-  MLIRArithmeticDialect
+  MLIRArithDialect
   MLIRAsyncDialect
   MLIRAffineUtils
   MLIRFuncDialect

lib/polygeist/Passes/CanonicalizeFor.cpp

@@ -9,7 +9,7 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "polygeist/Passes/Passes.h"
-#include <mlir/Dialect/Arithmetic/IR/Arithmetic.h>
+#include <mlir/Dialect/Arith/IR/Arith.h>
 
 using namespace mlir;
 using namespace mlir::scf;