Basic parfor-loop support in DaphneDSL #515

CMakeLists.txt

@@ -38,6 +38,16 @@ set(CMAKE_CXX_FLAGS_DEBUG="${CMAKE_CXX_FLAGS_DEBUG} -g -fno-omit-frame-pointer")
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3")
 set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -g -O3 -fno-omit-frame-pointer")
 
+option(ENABLE_PARALLEL_PARFOR "Enable PARALLEL_PARFOR" OFF)
+if(ENABLE_PARALLEL_PARFOR)
+    add_definitions(-DPARALLEL_PARFOR)
+endif()
+
+option(ENABLE_TIME_PARFOR "Enable TIME_PARFOR" OFF)
+if(ENABLE_TIME_PARFOR )
+    add_definitions(-DTIME_PARFOR)
+endif()
+
 # silence a warning about DEPFILE path transformations (used in LLVM)
 cmake_policy(SET CMP0116 OLD)
 

build.sh

@@ -62,6 +62,8 @@ function printHelp {
     echo "  --hdfs            Compile with support for HDFS"
     echo "  --io_uring        Compile with support for io_uring"
     echo "  --no-papi         Compile without support for PAPI"
+    echo "  --enable-parallel-parfor Compile with omp support for parfor statements"
+    echo "  --enable-time-parfor Run a timer during parfor execution and print the elapsed time"
 }
 
 #******************************************************************************
@@ -456,6 +458,8 @@ BUILD_IO_URING="-DUSE_IO_URING=OFF"
 BUILD_PAPI="-DUSE_PAPI=ON"
 WITH_DEPS=1
 WITH_SUBMODULE_UPDATE=1
+enableParallelParfor="OFF"
+enableTimeParfor="OFF"
 
 while [[ $# -gt 0 ]]; do
     key=$1
@@ -528,6 +532,14 @@ while [[ $# -gt 0 ]]; do
     -ns | --no-submodule-update)
         WITH_SUBMODULE_UPDATE=0
         ;;
+    --enable-parallel-parfor)
+        echo enabling parallelization for parfor
+        enableParallelParfor="ON"
+        ;;
+    --enable-time-parfor)
+        echo will run a timer during parfor execution and print the elapsed time
+        enableTimeParfor="ON"
+        ;;
       *)
         unknown_options="${unknown_options} ${key}"
         ;;
@@ -1109,6 +1121,9 @@ daphne_msg "Build Daphne"
 
 cmake -S "$projectRoot" -B "$daphneBuildDir" -G Ninja -DANTLR_VERSION="$antlrVersion" \
     -DCMAKE_PREFIX_PATH="$installPrefix" \
+    -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
+    -DENABLE_PARALLEL_PARFOR="$enableParallelParfor" \
+    -DENABLE_TIME_PARFOR="$enableTimeParfor" \
     $BUILD_CUDA $BUILD_FPGAOPENCL $BUILD_DEBUG $BUILD_MPI $BUILD_HDFS $BUILD_PAPI
 
 cmake --build "$daphneBuildDir" --target "$target"

src/compiler/execution/DaphneIrExecutor.cpp

@@ -84,16 +84,14 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
 
     // This flag is really useful to figure out why the lowering failed
     llvm::DebugFlag = userConfig_.debug_llvm;
-
+    context_.disableMultithreading();
     mlir::PassManager pm(&context_);
-    // TODO Enable the verifier for all passes where it is possible.
-    // Originally, it was only turned off for the
-    // SpecializeGenericFunctionsPass.
-    pm.enableVerifier(false);
+    pm.enableIRPrinting();
 
     if (userConfig_.explain_parsing)
         pm.addPass(mlir::daphne::createPrintIRPass("IR after parsing:"));
 
+    pm.addPass(mlir::daphne::createPrintIRPass("IR after parfor lowering:"));
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(mlir::createCSEPass());
     if (userConfig_.explain_parsing_simplified)
@@ -132,6 +130,7 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
         pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createInferencePass());
         // Simplify the IR.
         pm.addPass(mlir::createCanonicalizerPass());
+
         // Remove unused ops after simplifications.
         // TODO The CSE pass seems to eliminate only "one row" of dead code at a time, so we need it as many times as
         // the longest chain of ops we reduce; how to apply CSE until a fixpoint?
@@ -185,7 +184,6 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
         pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createProfilingPass());
 
     pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createInsertDaphneContextPass(userConfig_));
-
 #ifdef USE_CUDA
     if (userConfig_.use_cuda)
         pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createMarkCUDAOpsPass(userConfig_));
@@ -216,6 +214,7 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
     pm.addPass(mlir::createConvertSCFToCFPass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::LLVM::createRequestCWrappersPass());
     pm.addPass(mlir::daphne::createLowerToLLVMPass(userConfig_));
+    pm.addNestedPass<mlir::LLVM::LLVMFuncOp>(mlir::daphne::createLinkParForOutputPass());
     pm.addPass(mlir::createReconcileUnrealizedCastsPass());
     if (userConfig_.explain_llvm)
         pm.addPass(mlir::daphne::createPrintIRPass("IR after llvm lowering:"));

src/compiler/lowering/CMakeLists.txt

@@ -36,7 +36,7 @@ add_mlir_dialect_library(MLIRDaphneTransforms
     AggDimOpLowering.cpp
     TransposeOpLowering.cpp
     SparsityExploitationPass.cpp
-
+    LinkParForOutputPass.cpp
     DEPENDS
     MLIRDaphneOpsIncGen
     MLIRDaphneTransformsIncGen

src/compiler/lowering/LinkParForOutputPass.cpp

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2021 The DAPHNE Consortium
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ir/daphneir/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+#include <cstdio>
+#include <deque>
+#include <llvm/ADT/StringRef.h>
+#include <llvm/Support/Casting.h>
+#include <mlir/Analysis/SliceAnalysis.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <set>
+
+using namespace mlir;
+
+namespace {
+/**
+ * @brief Exploits the canonical form of ParFor output handling to rewrite the ParFor loop body function, 
+ * so that it conduct in-place updates on the shared buffer.  
+*/
+struct LinkParForOutputPass : public PassWrapper<LinkParForOutputPass, OperationPass<LLVM::LLVMFuncOp>> {
+    void runOnOperation() override {
+        LLVM::LLVMFuncOp func = getOperation();
+        llvm::StringRef fName = func.getSymName();
+
+        if (!fName.starts_with("parfor_body") || !func->hasAttr("parfor_inplace_rewrite_needed"))
+            return;
+
+        auto &blocks = func.getBody().getBlocks();
+        if (blocks.empty())
+            return;
+
+        // find all GEPOps that have the return arg as their base pointer
+        // output pointer is always passed first for kernels
+        auto funcOutArg = blocks.front().getArgument(0);
+
+        std::vector<LLVM::GEPOp> outGEPOps;
+        for (auto *user : funcOutArg.getUsers()) {
+            if (auto gep = dyn_cast<LLVM::GEPOp>(user)) {
+                if (gep.getBase() == funcOutArg)
+                    outGEPOps.push_back(gep);
+            }
+        }
+
+        // find stores to the GEPOps, these are what currently represents returns
+        std::vector<LLVM::StoreOp> returnStores;
+        for (auto gepOp : outGEPOps) {
+            for (auto user : gepOp->getUsers()) {
+                if (auto store = dyn_cast<LLVM::StoreOp>(user)) {
+                    if (store.getOperand(1) == gepOp) {
+                        returnStores.push_back(store);
+                    }
+                }
+            }
+        }
+        OpBuilder b(&getContext());
+
+        // rewire outputs of last kernel calls to the respective output of the function.
+        SetVector<Operation *> toErase = {};
+        for (auto store : returnStores) {
+            auto retVal = store->getOperand(0);
+            auto retValDef = retVal.getDefiningOp();
+            mlir::Value operand = retValDef->getOperand(0);
+            // load the output from the kernel
+            if (auto retValDef2 = retValDef->getOperand(0).getDefiningOp()) {
+                setInPlaceCalcKernelCall(retValDef2->getOperand(0).getDefiningOp(), store, b, &toErase);
+            } else if (auto blockArg = operand.dyn_cast<mlir::BlockArgument>()) {
+                setInPlaceCalcKernelCallInPrevBlocks(blockArg, store, b, &toErase);
+            }
+        }
+
+        // erase old operations after rewire
+        for (auto opToErase : toErase) {
+            opToErase->erase();
+        }
+        func->removeAttr("parfor_inplace_rewrite_needed");
+    }
+
+    /**
+     * @brief Traverses blocks backwards to determinate the `LLVM::LoadOp`, which loads the result of the last kernel
+     * call.
+     */
+    void setInPlaceCalcKernelCallInPrevBlocks(BlockArgument blockArg, mlir::Operation *store, OpBuilder b, SetVector<Operation *> *toErase) {
+        mlir::Block *parentBlock = blockArg.getOwner();
+        unsigned argIndex = blockArg.getArgNumber();
+
+        for (mlir::Block *pred : parentBlock->getPredecessors()) {
+            mlir::Operation &terminator = pred->back();
+
+            if (auto branchInterface = mlir::dyn_cast<mlir::BranchOpInterface>(terminator)) {
+                auto successors = terminator.getSuccessors();
+
+                for (unsigned succIdx = 0; succIdx < successors.size(); ++succIdx) {
+                    if (successors[succIdx] == parentBlock) {
+                        auto succOperands = branchInterface.getSuccessorOperands(succIdx);
+                        if (argIndex < succOperands.size()) {
+                            setInPlaceCalcKernelCall(succOperands[argIndex].getDefiningOp(), store, b, toErase);
+                        }
+                    }
+                }
+            } else if (terminator.getNumSuccessors() == 1 && terminator.getSuccessor(0) == parentBlock) {
+                auto succOperands = terminator.getOperands();
+                if (argIndex < succOperands.size()) {
+                    setInPlaceCalcKernelCall(succOperands[argIndex].getDefiningOp(), store, b, toErase);
+                }
+            }
+        }
+    }
+
+    /**
+     * @brief Replaces output argument of kernel CallOp with output argument of parfor body function to conduct in-place
+     * updates in shared buffer
+     */
+    void setInPlaceCalcKernelCall(mlir::Operation *load, mlir::Operation *store, OpBuilder b,
+                                  SetVector<Operation *> *toErase) {
+        auto ptr = load->getOperand(0);
+        mlir::Operation *lastUpdate = nullptr;
+        for (auto usr : ptr.getUsers()) {
+            if ((lastUpdate = llvm::dyn_cast<LLVM::CallOp>(usr))) {
+                break;
+            }
+        }
+        // Set the insertion point before lastUpdate
+        b.setInsertionPoint(lastUpdate);
+        // Get the defining operation of gep
+        auto gep = store->getOperand(1);
+        auto gepOp = gep.getDefiningOp();
+        auto offset = gepOp->getOperand(1).getDefiningOp();
+
+        // Insert a clone of gepOp at the new location (if moving, use move semantics if supported)
+        auto *clonedGepOp = gepOp->clone();
+        auto *clonedOffset = offset->clone();
+
+        clonedGepOp->setOperand(1, clonedOffset->getResult(0));
+
+        b.insert(clonedOffset);
+        b.insert(clonedGepOp);
+        // Update lastUpdate operand to use the result of the newly inserted gepOp
+        lastUpdate->setOperand(0, clonedGepOp->getResult(0));
+        // Load the in-place updated result
+        load->setOperand(0, clonedGepOp->getResult(0));
+        // Erase the old operations
+        toErase->insert(store);
+        toErase->insert(gepOp);
+        toErase->insert(offset);
+    }
+};
+
+} // end anonymous namespace
+
+std::unique_ptr<Pass> daphne::createLinkParForOutputPass() { return std::make_unique<LinkParForOutputPass>(); }