Add ISL support in polymer (#414)

.github/workflows/build.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   build:
-    name: Build ${{ matrix.os }} ${{ matrix.compiler }} polymer=${{ matrix.polymer }} pluto ${{ matrix.pluto-build }} mlir ${{ matrix.mlir-build }} polygeist ${{ matrix.polygeist-build }}
+    name: Build ${{ matrix.os }} ${{ matrix.compiler }} polymer=${{ matrix.polymer }} pluto ${{ matrix.pluto }} ${{ matrix.pluto-build }} mlir ${{ matrix.mlir-build }} polygeist ${{ matrix.polygeist-build }}
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -17,7 +17,8 @@ jobs:
         pluto-build: ["Release"] # Only Release supported currently
         mlir-build: ["Release"] #, "Debug"] # "RelWithDebInfo"
         polygeist-build: ["Release", "Debug"] #, "Debug"] # "RelWithDebInfo"
-        polymer: ["0", "1"]
+        polymer: ["1"]
+        pluto: ["1", "0"]
         os: [ubuntu-22.04]
         compiler: [gcc, clang]
         cxxcompiler: [g++, clang++]
@@ -26,6 +27,15 @@ jobs:
             cxxcompiler: clang++
           - compiler: clang
             cxxcompiler: g++
+        include:
+          - pluto-build: "Release"
+            mlir-build: "Release"
+            polygeist-build: "Release"
+            polymer: "0"
+            pluto: "0"
+            os: ubuntu-22.04
+            compiler: clang
+            cxxcompiler: clang++
 
     timeout-minutes: 360
     steps:
@@ -35,6 +45,11 @@ jobs:
         path: src
         submodules: 'recursive'
 
+    - name: define isl 1
+      id: defineisl
+      run: |
+        ( [ ${{matrix.pluto}} == 0 ] && echo "isl=1" || echo "isl=0" ) >> $GITHUB_OUTPUT
+
     - name: add dependencies
       run: sudo apt-get install -y llvm-15-tools ninja-build libbison-dev libtool libgmp-dev #libflex-dev cmake binutils-gold binutils binutils-dev ${{ matrix.compiler }} ${{ matrix.linker-pkg }}
 
@@ -46,13 +61,13 @@ jobs:
         key: ${{ matrix.pluto-build }}-${{ matrix.os }}-pluto-${{ matrix.polymer }}
 
     - name: Pluto build
-      if: steps.cache-pluto.outputs.cache-hit != 'true' && matrix.polymer == '1'
+      if: steps.cache-pluto.outputs.cache-hit != 'true' && matrix.polymer == '1' && matrix.pluto == '1'
       run: |
         mkdir -p pluto-build
         PATH="/usr/lib/llvm-15/bin/:$PATH" src/tools/polymer/build_polymer_deps.sh $(pwd)/pluto-build
         sleep 1s
     - uses: actions/cache/save@v3
-      if: steps.cache-pluto.outputs.cache-hit != 'true'
+      if: steps.cache-pluto.outputs.cache-hit != 'true' && matrix.polymer == '1' && matrix.pluto == '1'
       with:
         path: pluto-build
         key: ${{ matrix.pluto-build }}-${{ matrix.os }}-pluto-${{ matrix.polymer }}
@@ -66,14 +81,14 @@ jobs:
       uses: actions/cache@v3
       with:
         path: mlir-build
-        key: ${{ matrix.mlir-build }}-${{ matrix.os }}-mlir-${{ steps.getversion.outputs.version }}
+        key: ${{ matrix.mlir-build }}-${{ matrix.os }}-mlir-${{ steps.getversion.outputs.version }}-polly
 
     - name: MLIR build
       if: steps.cache-mlir.outputs.cache-hit != 'true'
       run: |
         mkdir mlir-build
         cd mlir-build
-        CYMBL=OFF cmake ../src/llvm-project/llvm -GNinja -DLLVM_ENABLE_PROJECTS="llvm;clang;mlir" -DCMAKE_BUILD_TYPE=${{ matrix.mlir-build }} -DCMAKE_C_COMPILER=/bin/clang -DCMAKE_CXX_COMPILER=/bin/clang++ -DCMAKE_ASM_COMPILER=/bin/clang -DCMAKE_CXX_FLAGS="-Wno-c++11-narrowing"
+        CYMBL=OFF cmake ../src/llvm-project/llvm -GNinja -DLLVM_ENABLE_PROJECTS="llvm;clang;mlir;polly" -DCMAKE_BUILD_TYPE=${{ matrix.mlir-build }} -DCMAKE_C_COMPILER=/bin/clang -DCMAKE_CXX_COMPILER=/bin/clang++ -DCMAKE_ASM_COMPILER=/bin/clang -DCMAKE_CXX_FLAGS="-Wno-c++11-narrowing"
         #cymbld & disown
         sleep 1
         CYMBL=OFF ninja
@@ -89,7 +104,7 @@ jobs:
       run: |
         cd build
         ls ../mlir-build/lib/cmake/clang
-        cmake ../src/ -GNinja -DMLIR_DIR=`pwd`/../mlir-build/lib/cmake/mlir -DLLVM_EXTERNAL_LIT=`pwd`/../mlir-build/bin/llvm-lit -DClang_DIR=`pwd`/../mlir-build/lib/cmake/clang -DCMAKE_BUILD_TYPE=${{ matrix.polygeist-build }} -DCMAKE_C_COMPILER=${{matrix.compiler}} -DCMAKE_CXX_COMPILER=${{matrix.cxxcompiler}} -DPOLYGEIST_ENABLE_POLYMER=${{ matrix.polymer }} -DPOLYMER_DEP_DIR=$(pwd)/../pluto-build/
+        cmake ../src/ -GNinja -DMLIR_DIR=`pwd`/../mlir-build/lib/cmake/mlir -DLLVM_EXTERNAL_LIT=`pwd`/../mlir-build/bin/llvm-lit -DClang_DIR=`pwd`/../mlir-build/lib/cmake/clang -DCMAKE_BUILD_TYPE=${{ matrix.polygeist-build }} -DCMAKE_C_COMPILER=${{matrix.compiler}} -DCMAKE_CXX_COMPILER=${{matrix.cxxcompiler}} -DPOLYGEIST_ENABLE_POLYMER=${{ matrix.polymer }} -DPOLYMER_DEP_DIR=$(pwd)/../pluto-build/ -DPOLYGEIST_POLYMER_ENABLE_PLUTO=${{matrix.pluto}} -DPOLYGEIST_POLYMER_ENABLE_ISL=${{ steps.defineisl.outputs.isl }}
 
     - name: test cgeist
       run: |
@@ -99,7 +114,7 @@ jobs:
         ninja check-cgeist
 
     - name: test polymer
-      if: matrix.polymer == '1'
+      if: matrix.polymer == '1' && matrix.pluto == '1'
       run: |
         cd build
         ninja

CMakeLists.txt

@@ -6,6 +6,8 @@ set(POLYGEIST_ENABLE_CUDA 0 CACHE BOOL "Enable CUDA frontend and backend")
 set(POLYGEIST_ENABLE_ROCM 0 CACHE BOOL "Enable ROCM backend")
 
 set(POLYGEIST_ENABLE_POLYMER 0 CACHE BOOL "Enable Polymer")
+set(POLYGEIST_POLYMER_ENABLE_ISL 0 CACHE BOOL "Enable Polymer isl")
+set(POLYGEIST_POLYMER_ENABLE_PLUTO 0 CACHE BOOL "Enable Polymer pluto")
 
 set(POLYMER_DEP_DIR "" CACHE STRING "Pluto build root directory")
 

README.md

@@ -1,23 +1,23 @@
-## Build instructions
+# Build instructions
 
-### Requirements 
+## Requirements 
 - Working C and C++ toolchains(compiler, linker)
 - cmake
 - make or ninja
 
-### 1. Clone Polygeist
+## 1. Clone Polygeist
 ```sh
 git clone --recursive https://github.com/llvm/Polygeist
 cd Polygeist
 ```
 
-### 2. Install LLVM, MLIR, Clang, and Polygeist
+## 2. Install LLVM, MLIR, Clang, and Polygeist
 
-#### Option 1: Using pre-built LLVM, MLIR, and Clang
+### Option 1: Using pre-built LLVM, MLIR, and Clang
 
 Polygeist can be built by providing paths to a pre-built MLIR and Clang toolchain.
 
-1. Build LLVM, MLIR, and Clang:
+#### 1. Build LLVM, MLIR, and Clang:
 ```sh
 mkdir llvm-project/build
 cd llvm-project/build
@@ -34,9 +34,11 @@ To enable compilation to cuda add `-DMLIR_ENABLE_CUDA_RUNNER=1` and remove `-DLL
 
 To enable the ROCM backend add `-DMLIR_ENABLE_ROCM_RUNNER=1` and remove `-DLLVM_TARGETS_TO_BUILD="host"` from the cmake arguments. (You may need to specify `-DHIP_CLANG_INCLUDE_PATH`, and/or `ROCM_PATH`)
 
+For ISL-enabled polymer, `polly` must be added to the `LLVM_ENABLE_PROJECTS` variable.
+
 For faster compilation we recommend using `-DLLVM_USE_LINKER=lld`.
 
-2. Build Polygeist:
+#### 2. Build Polygeist:
 ```sh
 mkdir build
 cd build
@@ -52,23 +54,31 @@ ninja check-polygeist-opt && ninja check-cgeist
 
 For faster compilation we recommend using `-DPOLYGEIST_USE_LINKER=lld`.
 
-    1. GPU backends
+##### GPU backends
 
 To enable the CUDA backend add `-DPOLYGEIST_ENABLE_CUDA=1`
 
 To enable the ROCM backend add `-DPOLYGEIST_ENABLE_ROCM=1`
 
-    2. Polymer
+##### Polymer
 
 To enable polymer, add `-DPOLYGEIST_ENABLE_POLYMER=1`
 
+There are two configurations of polymer that can be built - one with Pluto and one with ISL. 
+
+###### Pluto
+Add `-DPOLYGEIST_POLYMER_ENABLE_PLUTO=1`
 This will cause the cmake invokation to pull and build the dependencies for polymer. To specify a custom directory for the dependencies, specify `-DPOLYMER_DEP_DIR=<absolute-dir>`. The dependencies will be build using the `tools/polymer/build_polymer_deps.sh`.
 
-To run the polymer tests, use `ninja check-polymer`.
+To run the polymer pluto tests, use `ninja check-polymer`.
+
+###### ISL
 
+Add `-DPOLYGEIST_POLYMER_ENABLE_ISL=1`
+This requires an `llvm-project` build with `polly` enabled as a subproject.
 
 
-#### Option 2: Using unified LLVM, MLIR, Clang, and Polygeist build
+### Option 2: Using unified LLVM, MLIR, Clang, and Polygeist build
 
 Polygeist can also be built as an external LLVM project using [LLVM_EXTERNAL_PROJECTS](https://llvm.org/docs/CMake.html#llvm-related-variables).
 
@@ -90,7 +100,7 @@ ninja check-polygeist-opt && ninja check-cgeist
 `ninja check-polygeist-opt` runs the tests in `Polygeist/test/polygeist-opt`
 `ninja check-cgeist` runs the tests in `Polygeist/tools/cgeist/Test`
 
-## Citing Polygeist
+# Citing Polygeist
 
 If you use Polygeist, please consider citing the relevant publications:
 

include/polygeist/Passes/Passes.h

@@ -4,6 +4,7 @@
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "polygeist/Dialect.h"

include/polygeist/Passes/Passes.td

@@ -9,8 +9,15 @@ def PolyhedralOpt : Pass<"polyhedral-opt"> {
   let constructor = "mlir::polygeist::createPolyhedralOptPass()";
   let dependentDialects = [
     "scf::SCFDialect",
+    "func::FuncDialect",
     "arith::ArithDialect",
     "memref::MemRefDialect",
+    "affine::AffineDialect",
+    // This vector dialect is needed because lower-affine needs it. We launch
+    // a pass manager to use lower-affine from this pass itself, and that
+    // causes the multi-threaded mlir context to try to load the vector dialect
+    // which is unsupported. Instead we preload it here.
+    "vector::VectorDialect",
   ];
 }
 

lib/polygeist/Passes/PolyhedralOpt.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #ifdef POLYGEIST_ENABLE_POLYMER
 
 #include "mlir/IR/Visitors.h"
@@ -38,17 +39,28 @@
 using namespace mlir;
 using namespace polygeist;
 
+static llvm::cl::opt<std::string>
+    UsePolyhedralOptimizerCl("use-polyhedral-optimizer",
+                             llvm::cl::init("pluto"),
+                             llvm::cl::desc("pluto or islexternal"));
+
 namespace {
 
 #define POLYGEIST_OUTLINED_AFFINE_ATTR "polygeist.outlined_affine"
 
-static SmallVector<Operation *> findAffineRegions(Operation *root) {
-  SmallVector<Operation *> affineRegions;
+struct Scop {
+  Operation *begin;
+  Operation *end;
+};
+
+static SmallVector<Scop> findScops(Operation *root) {
+  SmallVector<Operation *> affineLoops;
+  SmallVector<Scop> scops;
   root->walk<mlir::WalkOrder::PreOrder>([&](Operation *loop) {
     if (!(isa<affine::AffineForOp>(loop) ||
           isa<affine::AffineParallelOp>(loop)))
       return;
-    if (!affineRegions.empty() && affineRegions.back()->isAncestor(loop))
+    if (!affineLoops.empty() && affineLoops.back()->isAncestor(loop))
       return;
 
     auto result = loop->walk<WalkOrder::PreOrder>([&](Operation *op) {
@@ -71,22 +83,34 @@ static SmallVector<Operation *> findAffineRegions(Operation *root) {
     if (!result.wasInterrupted()) {
       LLVM_DEBUG(llvm::dbgs() << DEBUG_LABEL << "Found affine region\n"
                               << *loop << "\n");
-      affineRegions.push_back(loop);
+      affineLoops.push_back(loop);
     }
   });
-  return affineRegions;
+  size_t size = affineLoops.size();
+  for (size_t i = 0; i < size; i++) {
+    Operation *loop = affineLoops[i];
+    Scop scop = {.begin = loop, .end = loop->getNextNode()};
+    while (i + 1 < size && scop.end == affineLoops[i + 1])
+      scop.end = affineLoops[++i]->getNextNode();
+    scops.push_back(scop);
+  }
+  return scops;
 }
 
 static FailureOr<func::FuncOp> outlineOp(RewriterBase &rewriter, Location loc,
-                                         Operation *op, StringRef funcName,
+                                         Scop scop, StringRef funcName,
                                          func::CallOp *callOp) {
   assert(!funcName.empty() && "funcName cannot be empty");
 
   OpBuilder::InsertionGuard g(rewriter);
-  rewriter.setInsertionPoint(op);
+  rewriter.setInsertionPoint(scop.begin);
   auto executeOp = rewriter.create<scf::ExecuteRegionOp>(loc, TypeRange());
   rewriter.createBlock(&executeOp.getRegion());
-  rewriter.clone(*op);
+  auto cur = scop.begin;
+  while (cur != scop.end) {
+    rewriter.clone(*cur);
+    cur = cur->getNextNode();
+  }
   rewriter.create<scf::YieldOp>(loc);
   auto ret = outlineSingleBlockRegion(rewriter, loc, executeOp.getRegion(),
                                       funcName, callOp);
@@ -96,15 +120,20 @@ static FailureOr<func::FuncOp> outlineOp(RewriterBase &rewriter, Location loc,
   }
   (*ret)->setAttr(POLYGEIST_OUTLINED_AFFINE_ATTR, rewriter.getUnitAttr());
   rewriter.eraseOp(executeOp.getRegion().front().getTerminator());
-  rewriter.inlineBlockBefore(&executeOp.getRegion().front(), op);
+  rewriter.inlineBlockBefore(&executeOp.getRegion().front(), scop.begin);
   rewriter.eraseOp(executeOp);
-  rewriter.eraseOp(op);
+  cur = scop.begin;
+  while (cur != scop.end) {
+    auto tmp = cur;
+    cur = cur->getNextNode();
+    rewriter.eraseOp(tmp);
+  }
   return ret;
 }
 
 static SmallVector<std::pair<func::FuncOp, func::CallOp>>
-outlineAffineRegions(Operation *root) {
-  auto affineRegions = findAffineRegions(root);
+outlineScops(Operation *root) {
+  auto scops = findScops(root);
   auto m = isa<ModuleOp>(root) ? cast<ModuleOp>(root)
                                : root->getParentOfType<ModuleOp>();
   auto loc = root->getLoc();
@@ -119,9 +148,9 @@ outlineAffineRegions(Operation *root) {
   };
   SmallVector<std::pair<func::FuncOp, func::CallOp>> funcs;
   IRRewriter rewriter(root->getContext());
-  for (Operation *op : affineRegions) {
+  for (Scop scop : scops) {
     func::CallOp callOp;
-    auto ret = outlineOp(rewriter, loc, op, getName(), &callOp);
+    auto ret = outlineOp(rewriter, loc, scop, getName(), &callOp);
     if (failed(ret)) {
       llvm::errs() << "Outlining affine region failed\n" << *root;
       abort();
@@ -179,11 +208,12 @@ void PolyhedralOptPass::runOnOperation() {
   ModuleOp m = cast<ModuleOp>(SymbolTable::getNearestSymbolTable(op));
   auto &context = *op->getContext();
 
-  auto funcOps = outlineAffineRegions(op);
+  auto funcOps = outlineScops(op);
   mlir::PassManager preTransformPm(&context);
   preTransformPm.addPass(createCanonicalizerPass());
 
-  AlwaysInlinerInterface interface(&getContext());
+  SmallVector<func::CallOp> toInline;
+
   IRRewriter b(&context);
   for (auto &pair : funcOps) {
     mlir::func::FuncOp f = pair.first;
@@ -205,18 +235,37 @@ void PolyhedralOptPass::runOnOperation() {
       return;
     }
     mlir::func::FuncOp g = nullptr;
-    if ((g = polymer::plutoTransform(f, b, ""))) {
+    if (UsePolyhedralOptimizerCl == "islexternal") {
+      g = polymer::islexternalTransform(f, b);
+    } else if (UsePolyhedralOptimizerCl == "pluto") {
+      g = polymer::plutoTransform(f, b, "");
+    }
+    if (g) {
       g.setPublic();
       g->setAttrs(f->getAttrs());
 
       g.setName(f.getName());
       f.erase();
     }
-    if (g && /*options.parallelize=*/true) {
-      polymer::plutoParallelize(g, b);
-    }
-    inlineAll(call);
+    // if (g && /*options.parallelize=*/true) {
+    //   polymer::plutoParallelize(g, b);
+    // }
+    toInline.push_back(call);
   }
+
+  mlir::PassManager lowerAffine(&context);
+  lowerAffine.addPass(createLowerAffinePass());
+  lowerAffine.addPass(createCanonicalizerPass());
+
+  // Conversion from ISL emits scf so we need to lower the statements before
+  // inlining them
+  if (UsePolyhedralOptimizerCl == "islexternal" && failed(lowerAffine.run(m))) {
+    signalPassFailure();
+    return;
+  }
+
+  for (auto call : toInline)
+    inlineAll(call);
   cleanupTempFuncs(m);
 }