Add lowering for loop.parallel to cfg.

[herhut]

Summary

This also removes the explicit pattern for loop.terminator to ensure
that the terminator is only erased if the parent op is rewritten.

Reductions are not yet supported.

[merge_guards_bot]

Unit tests: fail. 62155 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

Looks generally okay to me. I don't see the point of cloning before erasing, but can live with it pending a clean-up.

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

275

Ultra-nit: drop trivial braces.

276

Have you considered moving the region instead of cloning (inlineRegionBefore) and then calling replaceAllUsesWith to update induction variables?

mlir/test/Conversion/convert-to-cfg.mlir

174

std looks like leftover from the temporary state when loops where in the standard dialect, let's drop it here (and later clean up the tests above)

[herhut]

Review comments...

[herhut]

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

276

I did not want to rewrite the block arguments. In essence, I would need the equivalent of a signature conversion and that seemed involved. The issue here is that the block of the innermost loop only declares one induction variable.

mlir/test/Conversion/convert-to-cfg.mlir

174

I will clean it up and also move this test to LoopsToCfg as that test feels a bit misplaced.

[merge_guards_bot]

Unit tests: fail. 62155 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

let description = [{
     ```
    }];
 
   let arguments = (ins Variadic<Index>:$lowerBound,
                        Variadic<Index>:$upperBound,
                        Variadic<Index>:$step);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$body);
+
+  let extraClassDeclaration = [{
+    iterator_range<Block::args_iterator> getInductionVars() {
+      Block &block = body().front();
+      return {block.args_begin(), block.args_end()};
+    }
+  }];
 }
 
 def ReduceOp : Loop_Op<"reduce", [HasParent<"ParallelOp">]> {
   let summary = "reduce operation for parallel for";
   let description = [{
     "loop.reduce" is an operation occurring inside "loop.parallel" operations.
     It consists of one block with two arguments which have the same type as the
     operand of "loop.reduce".

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

 //===- ConvertLoopToStandard.cpp - ControlFlow to CFG conversion ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to convert loop.for, loop.if and loop.terminator
 // ops into standard CFG ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/Functional.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"

 //
 struct IfLowering : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(IfOp ifOp,
                                      PatternRewriter &rewriter) const override;
 };
 
-struct TerminatorLowering : public OpRewritePattern<TerminatorOp> {
-  using OpRewritePattern<TerminatorOp>::OpRewritePattern;
+struct ParallelLowering : public OpRewritePattern<mlir::loop::ParallelOp> {
+  using OpRewritePattern<mlir::loop::ParallelOp>::OpRewritePattern;
 
-  PatternMatchResult matchAndRewrite(TerminatorOp op,
-                                     PatternRewriter &rewriter) const override {
-    rewriter.eraseOp(op);
-    return matchSuccess();
-  }
+  PatternMatchResult matchAndRewrite(mlir::loop::ParallelOp parallelOp,
+                                     PatternRewriter &rewriter) const override;
 };
 } // namespace
 
 PatternMatchResult
 ForLowering::matchAndRewrite(ForOp forOp, PatternRewriter &rewriter) const {
   Location loc = forOp.getLoc();
 
   // Start by splitting the block containing the 'loop.for' into two parts.
auto *firstBodyBlock =
       rewriter.splitBlock(conditionBlock, conditionBlock->begin());
   auto *lastBodyBlock = &forOp.region().back();
   rewriter.inlineRegionBefore(forOp.region(), endBlock);
   auto iv = conditionBlock->getArgument(0);
 
   // Append the induction variable stepping logic to the last body block and
   // branch back to the condition block.  Construct an expression f :
   // (x -> x+step) and apply this expression to the induction variable.
+  rewriter.eraseOp(lastBodyBlock->getTerminator());
   rewriter.setInsertionPointToEnd(lastBodyBlock);
   auto step = forOp.step();
   auto stepped = rewriter.create<AddIOp>(loc, iv, step).getResult();
   if (!stepped)
     return matchFailure();
   rewriter.create<BranchOp>(loc, conditionBlock, stepped);
 
   // Compute loop bounds before branching to the condition.
IfLowering::matchAndRewrite(IfOp ifOp, PatternRewriter &rewriter) const {
   auto *condBlock = rewriter.getInsertionBlock();
   auto opPosition = rewriter.getInsertionPoint();
   auto *continueBlock = rewriter.splitBlock(condBlock, opPosition);
 
   // Move blocks from the "then" region to the region containing 'loop.if',
   // place it before the continuation block, and branch to it.
   auto &thenRegion = ifOp.thenRegion();
   auto *thenBlock = &thenRegion.front();
+  rewriter.eraseOp(thenRegion.back().getTerminator());
   rewriter.setInsertionPointToEnd(&thenRegion.back());
   rewriter.create<BranchOp>(loc, continueBlock);
   rewriter.inlineRegionBefore(thenRegion, continueBlock);
 
   // Move blocks from the "else" region (if present) to the region containing
   // 'loop.if', place it before the continuation block and branch to it.  It
   // will be placed after the "then" regions.
   auto *elseBlock = continueBlock;
   auto &elseRegion = ifOp.elseRegion();
   if (!elseRegion.empty()) {
     elseBlock = &elseRegion.front();
+    rewriter.eraseOp(elseRegion.back().getTerminator());
     rewriter.setInsertionPointToEnd(&elseRegion.back());
     rewriter.create<BranchOp>(loc, continueBlock);
     rewriter.inlineRegionBefore(elseRegion, continueBlock);
   }
 
   rewriter.setInsertionPointToEnd(condBlock);
   rewriter.create<CondBranchOp>(loc, ifOp.condition(), thenBlock,
                                 /*trueArgs=*/ArrayRef<Value>(), elseBlock,
                                 /*falseArgs=*/ArrayRef<Value>());
 
   // Ok, we're done!
   rewriter.eraseOp(ifOp);
   return matchSuccess();
 }
 
+PatternMatchResult
+ParallelLowering::matchAndRewrite(ParallelOp parallelOp,
+                                  PatternRewriter &rewriter) const {
+  Location loc = parallelOp.getLoc();
+  BlockAndValueMapping mapping;
+
+  if (parallelOp.getNumResults() != 0) {
+    // TODO: Implement lowering of parallelOp with reductions.
+    return matchFailure();
+  }
+
+  // For a parallel loop, we essentially need to create an n-dimensional loop
+  // nest. We do this by translating to loop.for ops and have those lowered in
+  // a further rewrite.
+  for (auto loop_operands :
+       llvm::zip(parallelOp.getInductionVars(), parallelOp.lowerBound(),
+                 parallelOp.upperBound(), parallelOp.step())) {
+    Value iv, lower, upper, step;
+    std::tie(iv, lower, upper, step) = loop_operands;
+    ForOp forOp = rewriter.create<ForOp>(loc, lower, upper, step);
+    mapping.map(iv, forOp.getInductionVar());
+    rewriter.setInsertionPointToStart(forOp.getBody());
+  }
+
+  // Now copy over the contents of the body.
+  for (auto &op : parallelOp.body().front().without_terminator())
+    rewriter.clone(op, mapping);
+
+  rewriter.eraseOp(parallelOp);
+
+  return matchSuccess();
+}
+
 void mlir::populateLoopToStdConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
-  patterns.insert<ForLowering, IfLowering, TerminatorLowering>(ctx);
+  patterns.insert<ForLowering, IfLowering, ParallelLowering>(ctx);
 }
 
 void LoopToStandardPass::runOnOperation() {
   OwningRewritePatternList patterns;
   populateLoopToStdConversionPatterns(patterns, &getContext());
   ConversionTarget target(getContext());
   target.addLegalDialect<StandardOpsDialect>();
   if (failed(applyPartialConversion(getOperation(), target, patterns)))

mlir/test/Conversion/convert-to-cfg.mlir

loop.if %arg3 {
         %c1_1 = constant 1 : index
       } else {
         %c1_2 = constant 1 : index
       }
     }
   }
   return
 }
+
+// CHECK-LABEL:   func @parallel_loop(
+// CHECK-SAME:                        [[VAL_0:%.*]]: index, [[VAL_1:%.*]]: index, [[VAL_2:%.*]]: index, [[VAL_3:%.*]]: index, [[VAL_4:%.*]]: index) {
+// CHECK:           [[VAL_5:%.*]] = constant 1 : index
+// CHECK:           br ^bb1([[VAL_0]] : index)
+// CHECK:         ^bb1([[VAL_6:%.*]]: index):
+// CHECK:           [[VAL_7:%.*]] = cmpi "slt", [[VAL_6]], [[VAL_2]] : index
+// CHECK:           cond_br [[VAL_7]], ^bb2, ^bb6
+// CHECK:         ^bb2:
+// CHECK:           br ^bb3([[VAL_1]] : index)
+// CHECK:         ^bb3([[VAL_8:%.*]]: index):
+// CHECK:           [[VAL_9:%.*]] = cmpi "slt", [[VAL_8]], [[VAL_3]] : index
+// CHECK:           cond_br [[VAL_9]], ^bb4, ^bb5
+// CHECK:         ^bb4:
+// CHECK:           [[VAL_10:%.*]] = constant 1 : index
+// CHECK:           [[VAL_11:%.*]] = addi [[VAL_8]], [[VAL_5]] : index
+// CHECK:           br ^bb3([[VAL_11]] : index)
+// CHECK:         ^bb5:
+// CHECK:           [[VAL_12:%.*]] = addi [[VAL_6]], [[VAL_4]] : index
+// CHECK:           br ^bb1([[VAL_12]] : index)
+// CHECK:         ^bb6:
+// CHECK:           return
+// CHECK:         }
+
+func @parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
+                        %arg3 : index, %arg4 : index) {
+  %step = constant 1 : index
+  loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
+                                          step (%arg4, %step) {
+    %c1 = constant 1 : index
+  }
+  return
+}