moved most iterateIE with 1 loop to forLoopIE new interface

Signed-off-by: Alexandre Eichenberger <[email protected]>

Author: AlexandreEichenberger

src/Conversion/ONNXToKrnl/NN/Normalization.cpp

@@ -941,17 +941,14 @@ struct GenericLayerNormaOpLowering : public OpConversionPattern<OP_TYPE> {
           invStdDevFlatMemRef);
     // Alloc mem for reductions (should be private if parallel)
     MemRefType tmpRedType = MemRefType::get({B, totVL}, elementType);
-    // Iterate over 1st dim by block
-    ValueRange loopDefs = create.krnl.defineLoops(1);
-    IndexExpr zero = LitIE(0);
-    ValueRange blockedLoopDefs = create.krnl.block(loopDefs[0], B);
-    Value blockedLoopDef = blockedLoopDefs[0];
+    // Iterate over 1st dim by block B.
+    bool useParallel = false;
     if (enableParallel) {
       int64_t parId;
       SmallVector<IndexExpr, 1> lb(1, LitIE(0)), ub(1, XFlatDims[0]);
       if (findSuitableParallelDimension(lb, ub, 0, 1, parId,
               /*min iter for going parallel*/ 4)) {
-        create.krnl.parallel(blockedLoopDef);
+        useParallel = true;
         onnxToKrnlParallelReport(op, true, 0, lb[0], ub[0], "in layer-norm");
       } else {
         onnxToKrnlParallelReport(
@@ -960,8 +957,7 @@ struct GenericLayerNormaOpLowering : public OpConversionPattern<OP_TYPE> {
     } else {
       onnxToKrnlParallelReport(op, false, -1, -1, "no parallel in layer norm");
     }
-    create.krnl.iterateIE({loopDefs[0]}, {blockedLoopDef}, {zero},
-        {XFlatDims[0]},
+    create.krnl.forLoopIE(LitIE(0), XFlatDims[0], /*step*/ B, useParallel,
         [&](const KrnlBuilder &ck, ValueRange blockedLoopIndices) {
           MDBuilder create(ck);
           IndexExprScope innerScope(ck);

src/Conversion/ONNXToKrnl/RNN/RNNBase.hpp

@@ -160,14 +160,13 @@ struct ONNXRNNOpLowering : public mlir::OpConversionPattern<RNNOp> {
 
     if (direction == FORWARD || direction == BIDIRECTIONAL) {
       IndexExprScope childScope(create.krnl);
-      mlir::ValueRange loopDef = create.krnl.defineLoops(1);
-      llvm::SmallVector<IndexExpr, 4> lbs(1, LitIE(0));
-      llvm::SmallVector<IndexExpr, 4> ubs;
+      IndexExpr lb = LitIE(0);
+      IndexExpr ub;
       if (!mlir::ShapedType::isDynamic(sequenceDimSize))
-        ubs.emplace_back(LitIE(sequenceDimSize));
+        ub = LitIE(sequenceDimSize);
       else
-        ubs.emplace_back(create.krnlIE.getShapeAsDim(X, 0));
-      create.krnl.iterateIE(loopDef, loopDef, lbs, ubs,
+        ub = create.krnlIE.getShapeAsDim(X, 0);
+      create.krnl.forLoopIE(lb, ub, /*step*/ 1, /*par*/ false,
           [&](const KrnlBuilder &createKrnl, mlir::ValueRange loopInd) {
             MathBuilder createMath(createKrnl);
             mlir::Value directionIV =
@@ -185,14 +184,13 @@ struct ONNXRNNOpLowering : public mlir::OpConversionPattern<RNNOp> {
 
     if (direction == REVERSE || direction == BIDIRECTIONAL) {
       IndexExprScope childScope(create.krnl);
-      mlir::ValueRange loopDef = create.krnl.defineLoops(1);
-      llvm::SmallVector<IndexExpr, 4> lbs(1, LitIE(0));
-      llvm::SmallVector<IndexExpr, 4> ubs;
+      IndexExpr lb = LitIE(0);
+      IndexExpr ub;
       if (!mlir::ShapedType::isDynamic(sequenceDimSize))
-        ubs.emplace_back(LitIE(sequenceDimSize));
+        ub = LitIE(sequenceDimSize);
       else
-        ubs.emplace_back(create.krnlIE.getShapeAsDim(X, 0));
-      create.krnl.iterateIE(loopDef, loopDef, lbs, ubs,
+        ub = create.krnlIE.getShapeAsDim(X, 0);
+      create.krnl.forLoopIE(lb, ub, /*step*/ 1, /*par*/ false,
           [&](const KrnlBuilder &ck, mlir::ValueRange loopInd) {
             MultiDialectBuilder<MemRefBuilder, MathBuilder> create(ck);
 

src/Conversion/ONNXToKrnl/Sequence/SequenceErase.cpp

@@ -64,12 +64,7 @@ struct ONNXSequenceEraseOpLowering
 
     // Copy the elements before the position
     KrnlBuilder createKrnl(rewriter, loc);
-    SmallVector<IndexExpr, 1> lbs;
-    lbs.emplace_back(LitIE(0));
-    SmallVector<IndexExpr, 1> ubs;
-    ubs.emplace_back(positionIE);
-    ValueRange firstLoopDef = createKrnl.defineLoops(1);
-    createKrnl.iterateIE(firstLoopDef, firstLoopDef, lbs, ubs,
+    createKrnl.forLoopIE(LitIE(0), positionIE, /*step*/ 1, /*par*/ false,
         [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
           Value element =
               createKrnl.load(adaptor.getInputSequence(), indicesLoopInd[0]);
@@ -78,12 +73,7 @@ struct ONNXSequenceEraseOpLowering
         });
 
     // Copy the elements after the position
-    SmallVector<IndexExpr, 1> lbs1;
-    lbs1.emplace_back(positionIE + 1);
-    SmallVector<IndexExpr, 1> ubs1;
-    ubs1.emplace_back(boundIE);
-    ValueRange secondLoopDef = createKrnl.defineLoops(1);
-    createKrnl.iterateIE(secondLoopDef, secondLoopDef, lbs1, ubs1,
+    createKrnl.forLoopIE(positionIE + 1, boundIE, /*step*/ 1, /*par*/ false,
         [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
           Value element =
               createKrnl.load(adaptor.getInputSequence(), indicesLoopInd[0]);

src/Conversion/ONNXToKrnl/Sequence/SequenceInsert.cpp

@@ -77,12 +77,7 @@ struct ONNXSequenceInsertOpLowering
       // compilation problem due to the unranked tensor even though
       // the loop will not be reached at runtime.
     } else {
-      SmallVector<IndexExpr, 1> lbs;
-      lbs.emplace_back(LitIE(0));
-      SmallVector<IndexExpr, 1> ubs;
-      ubs.emplace_back(positionIE);
-      ValueRange firstLoopDef = createKrnl.defineLoops(1);
-      createKrnl.iterateIE(firstLoopDef, firstLoopDef, lbs, ubs,
+      createKrnl.forLoopIE(LitIE(0), positionIE, /*step*/ 1, /*par*/ false,
           [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
             auto element =
                 createKrnl.load(adaptor.getInputSequence(), indicesLoopInd[0]);
@@ -91,12 +86,7 @@ struct ONNXSequenceInsertOpLowering
           });
 
       // Copy the elements after the position
-      SmallVector<IndexExpr, 1> lbs1;
-      lbs1.emplace_back(positionIE + 1);
-      SmallVector<IndexExpr, 1> ubs1;
-      ubs1.emplace_back(boundIE);
-      ValueRange secondLoopDef = createKrnl.defineLoops(1);
-      createKrnl.iterateIE(secondLoopDef, secondLoopDef, lbs1, ubs1,
+      createKrnl.forLoopIE(positionIE + 1, boundIE, /*step*/ 1, /*par*/ false,
           [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
             auto element =
                 createKrnl.load(adaptor.getInputSequence(), indicesLoopInd[0]);

src/Conversion/ONNXToKrnl/Tensor/Compress.cpp

@@ -60,8 +60,7 @@ struct ONNXCompressOpLowering : public OpConversionPattern<ONNXCompressOp> {
     // Now create a loop to iterate over all conditions.
     Value condMemRef = adaptor.getCondition();
     IndexExpr condShapeFirstRank = create.krnlIE.getShapeAsDim(condMemRef, 0);
-    ValueRange loopDef = create.krnl.defineLoops(1);
-    create.krnl.iterateIE(loopDef, loopDef, {zeroIE}, {condShapeFirstRank},
+    create.krnl.forLoopIE(zeroIE, condShapeFirstRank, /*step*/ 1, /*par*/ false,
         [&](const KrnlBuilder createKrnl, ValueRange loopInd) {
           MathBuilder createMath(createKrnl);
           // Load the condition
@@ -215,9 +214,8 @@ struct ONNXCompressOpLowering : public OpConversionPattern<ONNXCompressOp> {
         innerLbs.emplace_back(inputLbs[i]);
         innerUbs.emplace_back(inputUbs[i]);
       }
-      ValueRange axisLoopDef = create.krnl.defineLoops(1);
-      create.krnl.iterateIE(axisLoopDef, axisLoopDef, {inputLbs[axisValue]},
-          {inputUbs[axisValue]},
+      create.krnl.forLoopIE(inputLbs[axisValue], inputUbs[axisValue],
+          /*step*/ 1, /*par*/ false,
           [&](const KrnlBuilder createKrnl, ValueRange axisLoopInd) {
             MultiDialectBuilder<KrnlBuilder, MathBuilder, SCFBuilder> create(
                 createKrnl);

src/Conversion/ONNXToKrnl/Tensor/OneHot.cpp

@@ -62,7 +62,8 @@ struct ONNXOneHotOpLowering : public OpConversionPattern<ONNXOneHotOp> {
     create.krnlIE.getShapeAsDims(indices, indicesUbs);
     ValueRange indicesLoopDef = create.krnl.defineLoops(indicesRank);
     create.krnl.iterateIE(indicesLoopDef, indicesLoopDef, indicesLbs,
-        indicesUbs, [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
+        indicesUbs,
+        [&](const KrnlBuilder createKrnl, ValueRange indicesLoopInd) {
           // Loop for all input values.
           MathBuilder createMath(createKrnl);
           // Input val is allowed to be any integer/float. Read and convert to
@@ -89,8 +90,7 @@ struct ONNXOneHotOpLowering : public OpConversionPattern<ONNXOneHotOp> {
           Value onValueIndexVal = onValueIndex.getValue();
           // Now we have the index that is on, iterate over the depth values
           // along axis, and set the right one to the value on.
-          ValueRange depthLoopDef = createKrnl.defineLoops(1);
-          createKrnl.iterateIE(depthLoopDef, depthLoopDef, {zeroIE}, {depth},
+          createKrnl.forLoopIE(zeroIE, depth, /*step*/ 1, /*par*/ false,
               [&](const KrnlBuilder createBuilder, ValueRange depthLoopInd) {
                 MathBuilder createMath(createKrnl);
                 Value onCond = createMath.eq(depthLoopInd[0], onValueIndexVal);

src/Dialect/Krnl/DialectBuilder.cpp

@@ -286,55 +286,6 @@ void KrnlBuilder::simdReduce2DIE(IndexExpr lb, IndexExpr ub, int64_t VL,
       reductionBodyFn, postReductionBodyFn);
 }
 
-void KrnlBuilder::forExplicitlyParallelLoopIE(IndexExpr lb, IndexExpr ub,
-    int64_t stepModifier, IndexExpr numThreads, StringAttr procBind,
-    KrnlLoopBodyFn builderFn) const {
-  IndexExpr zero = LitIE(0);
-  if (numThreads.isLiteralAndIdenticalTo(1)) {
-    // Noop. Invoke function with (0, lb, ub).
-    SmallVector<Value, 4> params = {
-        zero.getValue(), lb.getValue(), ub.getValue()};
-    MultiDialectBuilder<KrnlBuilder> create(*this);
-    builderFn(create.krnl, params);
-    return;
-  }
-  if (numThreads.isLiteralAndIdenticalTo(-1)) {
-    // Dynamic, get value from OMP.
-    llvm_unreachable("not implemented yet"); // hi alex.
-  }
-  IndexExpr trip = ub - lb;
-  if (stepModifier > 1)
-    trip = trip.ceilDiv(stepModifier);
-  IndexExpr block = trip.ceilDiv(numThreads);
-  if (stepModifier > 1)
-    block = block * stepModifier;
-  // Create parallel loop with numThreads.
-  ValueRange originalLoopDef = defineLoops(1);
-  llvm::SmallVector<Value, 1> optLoopDef(1, originalLoopDef[0]);
-  parallel(optLoopDef[0], numThreads.getValue(), procBind);
-  iterateIE(originalLoopDef, optLoopDef, {lb}, {ub},
-      [&](const KrnlBuilder &kb, ValueRange loopInd) {
-        // Compute current LB/UB for this thread.
-        IndexExprScope currScope(kb);
-        IndexExpr tid = DimIE(loopInd[0]);
-        IndexExpr currLB = tid * SymIE(block);
-        IndexExpr currUB = currLB + SymIE(block);
-        currUB = IndexExpr::max(currUB, SymIE(ub));
-        SmallVector<Value, 4> params = {
-            tid.getValue(), currLB.getValue(), currUB.getValue()};
-        // Invoke function with (tid, currLB, currUB).
-        builderFn(kb, params);
-      });
-}
-
-void KrnlBuilder::forExplicitlyParallelLoopIE(IndexExpr lb, IndexExpr ub,
-    int64_t stepModifier, IndexExpr numThreads,
-    KrnlLoopBodyFn builderFn) const {
-  StringAttr procBind; // Empty == default, unspecified.
-  forExplicitlyParallelLoopIE(
-      lb, ub, stepModifier, numThreads, procBind, builderFn);
-}
-
 void KrnlBuilder::yield(ValueRange iterArgs) const {
   b().create<KrnlYieldOp>(loc(), iterArgs);
 }

src/Dialect/Krnl/DialectBuilder.hpp

@@ -106,22 +106,6 @@ struct KrnlBuilder : public DialectBuilder {
   void forLoopIE(IndexExpr lb, IndexExpr ub, int64_t step, bool useParallel,
       KrnlLoopBodyFn builderFn) const;
 
-  // Create a parallel loop iterating from tid=0 to tNum. When tNum== -1, the
-  // maximum number of thread is extracted at runtime (via is obtained from
-  // omp_get_num_threads). The parameters passed to the builder functions are
-  // (tid, currLB, currUB) where the currLB and currUB define the work assigned
-  // to a given thread tid. If follows the OpenMP static schedule, assigning
-  // roughly ceil((ub - lb)/tNum) iterations per thread. When stepModifier>1,
-  // assigned chunks are multiple of stepModifier.
-  // When tNum == 1, this is a nop, essentially calling the builder function
-  // with (0, lb, ub).
-  void forExplicitlyParallelLoopIE(IndexExpr lb, IndexExpr ub,
-      int64_t stepModifier, IndexExpr numThreads,
-      KrnlLoopBodyFn builderFn) const;
-  void forExplicitlyParallelLoopIE(IndexExpr lb, IndexExpr ub,
-      int64_t stepModifier, IndexExpr numThreads, mlir::StringAttr procBind,
-      KrnlLoopBodyFn builderFn) const;
-
   // Common simd loop interface (krnl/affine/scf).
   /*
      Iterate over a loop executing the loop body in SIMD mode (of vector length