[WIP] High-Level (Coase grain) parallelization with ONNXParallelOp and ONNXForkOp based on OpenMP

src/Compiler/CompilerPasses.cpp

@@ -186,13 +186,15 @@ void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
   if (enableInstrumentONNXSignature)
     pm.addNestedPass<func::FuncOp>(
         onnx_mlir::createInstrumentONNXSignaturePass());
-  pm.addPass(onnx_mlir::createLowerToKrnlPass(/*enableTiling*/ optLevel >= 3,
-      /*enableSIMD*/ optLevel >= 3 && !disableSimdOption, enableParallel,
-      /*opsToCall*/ opsForCall));
-  // An additional pass of canonicalization is helpful because lowering
-  // from ONNX dialect to Standard dialect exposes additional canonicalization
-  // opportunities.
-  pm.addPass(mlir::createCanonicalizerPass());
+  for (unsigned i = 0; i < 2; i++) {
+    pm.addPass(onnx_mlir::createLowerToKrnlPass(/*enableTiling*/ optLevel >= 3,
+        /*enableSIMD*/ optLevel >= 3 && !disableSimdOption, enableParallel,
+        /*opsToCall*/ opsForCall));
+    // An additional pass of canonicalization is helpful because lowering
+    // from ONNX dialect to Standard dialect exposes additional canonicalization
+    // opportunities.
+    pm.addPass(mlir::createCanonicalizerPass());
+  }
 }
 
 void addKrnlToAffinePasses(mlir::PassManager &pm) {

src/Conversion/KrnlToAffine/ConvertKrnlToAffine.cpp

@@ -17,6 +17,7 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"

src/Conversion/ONNXToKrnl/Additional/Parallel.cpp

@@ -0,0 +1,292 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===-------------------- Parallel.cpp - Lowering Parallel Op and Fork Op
+//---------------------===//
+//
+// Copyright 2019-2023 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the ONNX Parallel and Fork Operators to Krnl dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
+#include "src/Dialect/Krnl/DialectBuilder.hpp"
+
+#include <llvm/Support/Debug.h>
+
+#define DEBUG_TYPE "lowering-parallelop-to-krnl"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+//===----------------------------------------------------------------------===//
+// Helper function
+// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
+// properly dominates `b` and `b` is not inside `a`.
+// Reference: llvm-project/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
+//===----------------------------------------------------------------------===//
+
+static bool happensBefore(Operation *a, Operation *b) {
+  do {
+    if (a->isProperAncestor(b))
+      return false;
+    if (Operation *bAncestor = a->getBlock()->findAncestorOpInBlock(*b)) {
+      return a->isBeforeInBlock(bAncestor);
+    }
+  } while ((a = a->getParentOp()));
+  return false;
+}
+
+void moveAllocOpOperands(SmallVector<Operation *, 4> &opsToMove,
+    SmallVector<Operation *, 4> &globalOpsToMove, Operation *returnValOp,
+    Operation *parentOp) {
+  if (opsToMove.size() == 0)
+    return;
+
+  SmallVector<Operation *, 4> nextOpsToMove;
+  for (Operation *op : opsToMove) {
+    LLVM_DEBUG(llvm::dbgs() << "@@START opsToMove op: = " << *op << "\n");
+    // Added the op in ops list to move if it is still not added.
+    if (llvm::find(globalOpsToMove, op) == globalOpsToMove.end()) {
+      globalOpsToMove.push_back(op);
+      LLVM_DEBUG(llvm::dbgs() << "Added in opsToMove : " << *op << "\n");
+    }
+
+    Region &parentOpRegion = parentOp->getRegions().front();
+    Block &parentOpBlock = parentOpRegion.getBlocks().front();
+
+    // AllocOp: If allocated value is used in KrnlStoreOp, the KrnlStoreOp need
+    // to be move.
+    if (op != returnValOp) {
+      if (auto allocOp = dyn_cast<memref::AllocOp>(op)) {
+        for (Operation *user : allocOp.getResult().getUsers()) {
+          if (auto krnlStoreOp = dyn_cast<KrnlStoreOp>(user)) {
+            if (user->getBlock() == &parentOpBlock) {
+              if ((llvm::find(nextOpsToMove, user) == nextOpsToMove.end()) and
+                  (llvm::find(globalOpsToMove, user) ==
+                      globalOpsToMove.end())) {
+                LLVM_DEBUG(llvm::dbgs()
+                           << "Added in nextOpsTomove (Single KrnlStore) = "
+                           << *user << "\n");
+                nextOpsToMove.push_back(user);
+              }
+            } else {
+              if ((llvm::find(nextOpsToMove, user->getParentOp()) ==
+                      nextOpsToMove.end()) and
+                  (llvm::find(globalOpsToMove, user->getParentOp()) ==
+                      globalOpsToMove.end())) {
+                LLVM_DEBUG(llvm::dbgs()
+                           << "Added in nextOpsTomove (KrnlIterateOp): "
+                           << *(user->getParentOp()) << "\n");
+                nextOpsToMove.push_back(user->getParentOp());
+              }
+            }
+          }
+        }
+      }
+    }
+    // KrnlIterateOp: Operations in the region of KrnlIterateOp are already
+    // added in opsToMove. So, add operands of operations in the region of
+    // KrnlIterateOp.
+    if (auto iterateOp = dyn_cast<KrnlIterateOp>(op)) {
+      Block &iterationBlock = iterateOp.getBodyRegion().front();
+      for (Operation &iop : iterationBlock.getOperations()) {
+        LLVM_DEBUG(llvm::dbgs() << "Ops in krnlIterateOp: " << *(&iop) << "\n");
+        for (unsigned i = 0; i < iop.getNumOperands(); ++i) {
+          Value oprd = iop.getOperand(i);
+          if (isa<BlockArgument>(oprd))
+            continue;
+          Operation *oprdDefOp = oprd.getDefiningOp();
+          if (oprdDefOp->getBlock() != &iterationBlock and
+              oprdDefOp->getBlock() == &parentOpBlock) {
+            if ((llvm::find(nextOpsToMove, oprdDefOp) ==
+                    nextOpsToMove.end()) and
+                (llvm::find(globalOpsToMove, oprdDefOp) ==
+                    globalOpsToMove.end())) {
+
+              LLVM_DEBUG(llvm::dbgs()
+                         << "Added in nextOpsTomove operand in KrnlIterateOp: "
+                         << *oprdDefOp << "\n");
+              nextOpsToMove.push_back(oprdDefOp);
+            }
+          }
+        }
+      }
+    }
+
+    // Check if operands need to be moved. Need to move if defining op for the
+    // operand exists in block in parentOp.
+    for (unsigned i = 0; i < op->getNumOperands(); ++i) {
+      Value oprd = op->getOperand(i);
+      if (isa<BlockArgument>(oprd))
+        continue;
+      Operation *oprdDefOp = oprd.getDefiningOp();
+      if (oprdDefOp->getBlock() == &parentOpBlock) {
+        if ((llvm::find(nextOpsToMove, oprdDefOp) == nextOpsToMove.end()) and
+            (llvm::find(globalOpsToMove, oprdDefOp) == globalOpsToMove.end())) {
+          LLVM_DEBUG(llvm::dbgs() << "Added in nextOpsTomove operand for op: "
+                                  << i << " = " << *oprdDefOp << "\n");
+          nextOpsToMove.push_back(oprdDefOp);
+        }
+      }
+    }
+    LLVM_DEBUG(llvm::dbgs() << "@@END\n");
+  }
+  // Check if operands of operandDefOp need to be moved recursively
+  moveAllocOpOperands(nextOpsToMove, globalOpsToMove, returnValOp, parentOp);
+}
+
+LogicalResult moveAllocOpBeforeAndReplaceAllUses(
+    ConversionPatternRewriter &rewriter, Operation *op, Operation *yieldOp) {
+  SmallVector<Operation *, 4> globalOpsToMove;
+  for (unsigned ii = 0; ii < yieldOp->getNumOperands(); ++ii) {
+    // Check the return value of the block to check if operations in the
+    // block is already lowered to memref-level IR such as KrnlIR. Assume
+    // the block is still not lowered if the return value is still Tensor
+    // type. Actual return value of ONNXYieldOp is conveted into tensor by
+    // unrealized_conversion_cast. So, check the operand of previous
+    // operation.
+    Value returnVal = yieldOp->getOperands()[ii];
+    if (isa<UnrealizedConversionCastOp>(returnVal.getDefiningOp()))
+      returnVal = returnVal.getDefiningOp()->getOperands()[0];
+    if (isa<TensorType>(returnVal.getType()))
+      return failure();
+
+    // Move allocOps for results before op
+    Operation *allocOpForReturnVal = returnVal.getDefiningOp();
+    SmallVector<Operation *, 4> opsToMove;
+    opsToMove.push_back(allocOpForReturnVal);
+    moveAllocOpOperands(opsToMove, globalOpsToMove, allocOpForReturnVal, op);
+    rewriter.replaceAllUsesWith(
+        op->getResults()[ii], allocOpForReturnVal->getResult(0));
+  }
+
+  llvm::sort(globalOpsToMove,
+      [](Operation *a, Operation *b) { return !happensBefore(a, b); });
+  Operation *justMovedOp = op;
+  for (Operation *gop : globalOpsToMove) {
+    gop->moveBefore(justMovedOp);
+    justMovedOp = gop;
+  }
+  return success();
+}
+
+void insertRegionInIterateOp(
+    ConversionPatternRewriter &rewriter, Location &loc, Block &block) {
+  for (auto iterateOp : block.getOps<KrnlIterateOp>()) {
+    // Currently KrnlRegionOp does not support input and output to the region.
+    // So, if KrnlIterateOp has output (by krnl.yield), can't use KrnlRegionOp.
+    if (iterateOp.getNumResults() == 0) {
+      KrnlRegionOp regionOp = rewriter.create<KrnlRegionOp>(loc);
+      Block &regionBlock = regionOp.getBodyRegion().front();
+      Block &iterateBlock = iterateOp.getRegion().back();
+      insertRegionInIterateOp(rewriter, loc, iterateBlock); // recursive call
+      rewriter.eraseOp(iterateBlock.getTerminator());
+      regionBlock.getOperations().splice(
+          regionBlock.end(), iterateBlock.getOperations());
+      rewriter.setInsertionPointToStart(&iterateBlock);
+      KrnlYieldOp krnlYieldOp = rewriter.create<KrnlYieldOp>(loc);
+      rewriter.moveOpBefore(regionOp, krnlYieldOp);
+    }
+  }
+}
+
+struct ONNXParallelOpLowering : public OpConversionPattern<ONNXParallelOp> {
+  explicit ONNXParallelOpLowering(
+      TypeConverter &typeConverter, MLIRContext *ctx)
+      : OpConversionPattern(typeConverter, ctx) {}
+
+  LogicalResult matchAndRewrite(ONNXParallelOp parallelOp,
+      ONNXParallelOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const final {
+    Operation *op = parallelOp.getOperation();
+    Location loc = ONNXLoc<ONNXParallelOp>(op);
+    IndexExprScope ieScope(&rewriter, loc);
+    MultiDialectBuilder<KrnlBuilder, IndexExprBuilderForKrnl, MemRefBuilder,
+        MathBuilder>
+        create(rewriter, loc);
+
+    auto onnxParallelOp = dyn_cast<ONNXParallelOp>(op);
+    // Get the parallel region.
+    Region &parallelBody = onnxParallelOp.getBody();
+    // Make sure the region has only one block.
+    if (!parallelBody.hasOneBlock())
+      return success();
+    // Get YieldOp of the body block.
+    Block &bodyBlock = parallelBody.front();
+    Operation *yieldOp = bodyBlock.getTerminator();
+    if (!isa<ONNXYieldOp>(yieldOp))
+      return failure();
+
+    // Move alloc ops included in ForkOps
+    SmallVector<ONNXForkOp, 4> forkOps;
+    for (Operation &bOp : bodyBlock.getOperations()) {
+      if (auto forkOp = dyn_cast<ONNXForkOp>(bOp)) {
+        forkOps.push_back(forkOp);
+        Operation *forkYieldOp = forkOp.getBody().front().getTerminator();
+        if (!isa<ONNXYieldOp>(forkYieldOp))
+          return failure();
+
+        if (failed(moveAllocOpBeforeAndReplaceAllUses(
+                rewriter, &bOp, forkYieldOp)))
+          return failure();
+      }
+    }
+
+    // Move allocOp included in ParallelOp
+    if (failed(moveAllocOpBeforeAndReplaceAllUses(rewriter, op, yieldOp)))
+      return failure();
+
+    // Create KrnlIterateOp and replace ParallelOp with it.
+    rewriter.setInsertionPoint(op);
+    std::vector<Value> loop;
+    defineLoops(rewriter, loc, loop, 1);
+    krnl::KrnlIterateOperandPack pack(rewriter, loop);
+    pack.pushConstantBound(0);
+    pack.pushConstantBound(forkOps.size());
+    KrnlBuilder createKrnl(rewriter, loc);
+    createKrnl.parallel(loop);
+    KrnlIterateOp iterateOp = createKrnl.iterate(pack);
+    Block &iterationBlock = iterateOp.getBodyRegion().back();
+    rewriter.setInsertionPointToStart(&iterationBlock);
+    ValueRange indices = createKrnl.getInductionVarValue({loop[0]});
+    rewriter.eraseOp(yieldOp);
+    rewriter.inlineBlockBefore(&bodyBlock, iterationBlock.getTerminator());
+
+    // Create SCFIfOp and replace ForkOp with it.
+    int64_t id = 0;
+    for (auto forkOp : forkOps) {
+      rewriter.setInsertionPoint(forkOp);
+      // Insert scf::IfOp
+      Value forkId = create.math.constantIndex(id);
+      Value eq = create.math.eq(forkId, indices[0]);
+      scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, eq, /*else=*/false);
+      Block &ifBlock = ifOp.getThenRegion().back();
+      rewriter.setInsertionPointToStart(&ifBlock);
+      // Insert KrnlRegionOp in every KrnlIterateOps. This needs to avoid
+      // errors in convertKrnlToAffinePass.
+      Block &forkBlock = forkOp.getRegion().back();
+      insertRegionInIterateOp(rewriter, loc, forkBlock);
+      Operation *forkYieldOp = forkBlock.getTerminator();
+      rewriter.eraseOp(forkYieldOp);
+      rewriter.inlineBlockBefore(&forkBlock, ifBlock.getTerminator());
+      rewriter.eraseOp(forkOp);
+      id++;
+    }
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+void populateLoweringONNXParallelOpPattern(RewritePatternSet &patterns,
+    TypeConverter &typeConverter, MLIRContext *ctx) {
+  patterns.insert<ONNXParallelOpLowering>(typeConverter, ctx);
+}
+
+} // namespace onnx_mlir

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -8,6 +8,7 @@ add_onnx_mlir_library(OMONNXToKrnl
   Additional/Custom.cpp
   Additional/LayoutTransform.cpp
   Additional/ShapeTransform.cpp
+  Additional/Parallel.cpp
   ControlFlow/If.cpp
   ControlFlow/Loop.cpp
   ControlFlow/Scan.cpp

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -286,6 +286,7 @@ void populateONNXToKrnlConversionPattern(RewritePatternSet &patterns,
   populateLoweringONNXCustomOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXLayoutTransformOpPattern(patterns, typeConverter, ctx, enableParallel);
   populateLoweringONNXShapeTransformOpPattern(patterns, typeConverter, ctx);
+  populateLoweringONNXParallelOpPattern(patterns, typeConverter, ctx);
   // clang-format on
 }
 

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp

@@ -24,6 +24,7 @@
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -470,6 +471,9 @@ void populateLoweringONNXShapeTransformOpPattern(
 void populateLoweringONNXCustomOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 
+void populateLoweringONNXParallelOpPattern(
+    mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
+
 // Utilities for generating krnl.call for ONNX Ops
 
 // Create allocate based on COMPUTED shapeHelper.