Parallel tune: transformations for Roberta with BS=1 (onnx#2840)

Signed-off-by: Alexandre Eichenberger <[email protected]>
Co-authored-by: Tung D. Le <[email protected]>

src/Conversion/ONNXToKrnl/Math/Reduction.cpp

@@ -778,7 +778,6 @@ struct ONNXReductionOpLowering : public OpConversionPattern<ONNXReductionOp> {
         int64_t parId;
         if (findSuitableParallelDimension(lbs3, ubs3, 0, 1, parId,
                 /*min iter for going parallel*/ 4)) {
-
           create.krnl.parallel(loop3Def[0]);
           onnxToKrnlParallelReport(
               op, true, 0, lbs3[0], ubs3[0], "reduction scalar mean");

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp

@@ -617,17 +617,23 @@ bool hasNonIdentityLayout(ValueRange operands) {
 // Support functions for parallel region.
 //===----------------------------------------------------------------------===//
 
-// Return the outermost loop within [firstDim, lastDim) for which (ub-lb) >
-// minSize. Runtime dimensions are assumed to satisfy the size requirement by
-// definition. If found one, it is parDim and the function returns true.
+// Return the outermost loop within [firstInclusiveDim, lastExclusiveDim) for
+// which (ub-lb) > minSize. Runtime dimensions are assumed to satisfy the size
+// requirement by definition. If found one, it is parDim and the function
+// returns true.
 
 bool findSuitableParallelDimension(llvm::SmallVectorImpl<IndexExpr> &lb,
-    llvm::SmallVectorImpl<IndexExpr> &ub, int64_t firstDim, int64_t lastDim,
-    int64_t &parDim, int64_t minSize) {
-  for (int64_t i = firstDim; i < lastDim; ++i) {
+    llvm::SmallVectorImpl<IndexExpr> &ub, int64_t firstInclusiveDim,
+    int64_t lastExclusiveDim, int64_t &parDim, int64_t minSize) {
+  assert(lb.size() == ub.size() && "expected identical ranks for lb/ub");
+  if (firstInclusiveDim < 0)
+    firstInclusiveDim = 0;
+  if (lastExclusiveDim > (int64_t)lb.size())
+    lastExclusiveDim = lb.size();
+  for (int64_t i = firstInclusiveDim; i < lastExclusiveDim; ++i) {
     IndexExpr tripCount = ub[i] - lb[i];
     if (!tripCount.isLiteral() || tripCount.getLiteral() >= minSize) {
-      // Got one
+      // Got one.
       parDim = i;
       return true;
     }

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp

@@ -592,8 +592,8 @@ bool hasNonIdentityLayout(mlir::ValueRange operands);
 // minSize. Runtime dimensions are assumed to satisfy the size requirement by
 // definition. If found one, it is parDim and the function returns true.
 bool findSuitableParallelDimension(llvm::SmallVectorImpl<IndexExpr> &lb,
-    llvm::SmallVectorImpl<IndexExpr> &ub, int64_t firstDim /*inclusive*/,
-    int64_t lastDim /*exclusive*/, int64_t &parDim, int64_t minSize = 4);
+    llvm::SmallVectorImpl<IndexExpr> &ub, int64_t firstInclusiveDim,
+    int64_t lastExclusiveDim, int64_t &parDim, int64_t minSize = 4);
 
 //===----------------------------------------------------------------------===//
 // Support functions for reporting.

src/Conversion/ONNXToKrnl/Tensor/Transpose.cpp

@@ -150,11 +150,10 @@ struct ONNXTransposeOpLowering : public OpConversionPattern<ONNXTransposeOp> {
     if (enableParallel) {
       int64_t parId;
       // TODO: consider flattening the outer dims, or along inner dims.
-      if (findSuitableParallelDimension(lbs, ubs, 0, 1, parId, 8)) {
-        assert(parId == 0 && "only outermost at this time");
-        create->krnl.parallel(loopDef[0]);
+      if (findSuitableParallelDimension(lbs, ubs, 0, 2, parId, 8)) {
+        create->krnl.parallel(loopDef[parId]);
         onnxToKrnlParallelReport(
-            op, true, 0, lbs[0], ubs[0], "scalar transpose");
+            op, true, parId, lbs[parId], ubs[parId], "scalar transpose");
       } else {
         onnxToKrnlParallelReport(
             op, false, -1, -1, "no dim with enough work in scalar transpose");
@@ -224,12 +223,12 @@ struct ONNXTransposeOpLowering : public OpConversionPattern<ONNXTransposeOp> {
     SmallVector<IndexExpr, 4> lbs(outerRank, LiteralIndexExpr(0));
     if (enableParallel) {
       int64_t parId;
-      // TODO: consider flattening the outer dims, or along inner dims.
-      if (findSuitableParallelDimension(lbs, inUBs, 0, 1, parId, 8)) {
-        assert(parId == 0 && "only outermost at this time");
-        create->krnl.parallel(loopDef[0]);
+      // Note that if there is only 1 dim, lastExclusiveDim is automatically
+      // reduced to 1 in the findSuitableParallelDimension call.
+      if (findSuitableParallelDimension(lbs, inUBs, 0, 2, parId, 8)) {
+        create->krnl.parallel(loopDef[parId]);
         onnxToKrnlParallelReport(
-            op, true, 0, lbs[0], inUBs[0], "block transpose");
+            op, true, parId, lbs[parId], inUBs[parId], "block transpose");
       } else {
         onnxToKrnlParallelReport(
             op, false, -1, -1, "no dim with enough work in block transpose");