Add canonicalizer removing rank 0 affine.parallel

Author: tzerrell

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -587,6 +587,8 @@ def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> {
     static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
     static StringRef getStepsAttrName() { return "steps"; }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def AffinePrefetchOp : Affine_Op<"prefetch"> {

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -2614,6 +2614,33 @@ static LogicalResult verify(AffineVectorStoreOp op) {
   return success();
 }
 
+namespace {
+/// This pattern removes affine.parallel ops with no induction variables
+struct AffineParallelRank0LoopRemover
+    : public OpRewritePattern<AffineParallelOp> {
+  using OpRewritePattern<AffineParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check that there are no induction variables
+    if (op.lowerBoundsMap().getNumResults() != 0)
+      return failure();
+    // Remove the affine.parallel wrapper, retain the body in the same location
+    auto &parentOps = rewriter.getInsertionBlock()->getOperations();
+    auto &parallelBodyOps = op.region().front().getOperations();
+    parentOps.splice(mlir::Block::iterator(op), parallelBodyOps,
+                     parallelBodyOps.begin(), std::prev(parallelBodyOps.end()));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+} // end anonymous namespace
+
+void AffineParallelOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<AffineParallelRank0LoopRemover>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Affine/canonicalize.mlir

@@ -604,3 +604,18 @@ func @drop_duplicate_bounds(%N : index) {
   }
   return
 }
+
+// -----
+
+// CHECK: func @remove_rank0_affine_parallel(%[[OUT:.*]]: memref<f32>)
+func @remove_rank0_affine_parallel(%out: memref<f32>) {
+  // CHECK-NEXT: %[[CST:.*]] = constant
+  %cst = constant 0.0 : f32
+  // CHECK-NEXT: affine.store %[[CST]], %[[OUT]][] : memref<f32>
+  affine.parallel () = () to () {
+    affine.parallel () = () to () {
+      affine.store %cst, %out[] : memref<f32>
+    }
+  }
+  return
+}