make Affine parallel and yield ops MemRefsNormalizable

Affine parallel ops may contain and yield results from MemRefsNormalizable ops in the loop body.  Thus, both affine.parallel and affine.yield should have the MemRefsNormalizable trait.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D96821

Author: adstraw

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

@@ -592,7 +592,7 @@ def AffineMaxOp : AffineMinMaxOpBase<"max", [NoSideEffect]> {
 
 def AffineParallelOp : Affine_Op<"parallel",
     [ImplicitAffineTerminator, RecursiveSideEffects,
-     DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
+     DeclareOpInterfaceMethods<LoopLikeOpInterface>, MemRefsNormalizable]> {
   let summary = "multi-index parallel band operation";
   let description = [{
     The "affine.parallel" operation represents a hyper-rectangular affine
@@ -842,7 +842,8 @@ def AffineStoreOp : AffineStoreOpBase<"store"> {
   let hasFolder = 1;
 }
 
-def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator, ReturnLike]> {
+def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator, ReturnLike,
+    MemRefsNormalizable]> {
   let summary = "Yield values to parent operation";
   let description = [{
     "affine.yield" yields zero or more SSA values from an affine op region and

mlir/lib/Transforms/NormalizeMemRefs.cpp

@@ -512,6 +512,10 @@ Operation *NormalizeMemRefs::createOpResultsNormalized(FuncOp funcOp,
   // affine map, `oldOp` is returned without modification.
   if (resultTypeNormalized) {
     OpBuilder bb(oldOp);
+    for (auto &oldRegion : oldOp->getRegions()) {
+      Region *newRegion = result.addRegion();
+      newRegion->takeBody(oldRegion);
+    }
     return bb.createOperation(result);
   } else
     return oldOp;

mlir/test/Transforms/normalize-memrefs.mlir

@@ -319,3 +319,16 @@ func @use_value_of_external(%A: memref<16xf64, #tile>, %B: f64) -> (memref<8xf64
 }
 // CHECK: %[[res:[0-9]+]] = call @external_func_B(%[[A]], %[[B]]) : (memref<4x4xf64>, f64) -> memref<2x4xf64>
 // CHECK: return %{{.*}} : memref<2x4xf64>
+
+// CHECK-LABEL: func @affine_parallel_norm
+func @affine_parallel_norm() ->  memref<8xf32, #tile> {
+  %c = constant 23.0 : f32
+  %a = alloc() : memref<8xf32, #tile>
+  // CHECK: affine.parallel (%{{.*}}) = (0) to (8) reduce ("assign") -> (memref<2x4xf32>)
+  %1 = affine.parallel (%i) = (0) to (8) reduce ("assign") ->  memref<8xf32, #tile> {
+    affine.store %c, %a[%i] : memref<8xf32, #tile>
+    // CHECK: affine.yield %{{.*}} : memref<2x4xf32>
+    affine.yield %a : memref<8xf32, #tile>
+  }
+  return %1 : memref<8xf32, #tile>
+}