Add structural attr to AffineParallelOps

Author: tzerrell

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

@@ -549,6 +549,7 @@ def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> {
      AffineMapAttr:$lowerBoundsMap,
      AffineMapAttr:$upperBoundsMap,
      I64ArrayAttr:$steps,
+     OptionalAttr<UnitAttr>:$structural,
      Variadic<Index>:$mapOperands);
   let regions = (region SizedRegion<1>:$region);
 

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

@@ -2622,8 +2622,8 @@ struct AffineParallelRank0LoopRemover
 
   LogicalResult matchAndRewrite(AffineParallelOp op,
                                 PatternRewriter &rewriter) const override {
-    // Check that there are no induction variables
-    if (op.lowerBoundsMap().getNumResults() != 0)
+    // Check that there are no induction variables & op isn't structural
+    if (op.lowerBoundsMap().getNumResults() != 0 || op.getAttr("structural"))
       return failure();
     // Remove the affine.parallel wrapper, retain the body in the same location
     auto &parentOps = rewriter.getInsertionBlock()->getOperations();

mlir/test/Dialect/Affine/canonicalize.mlir

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