make Affine parallel and yield ops MemRefsNormalizable

[adstraw]

Summary

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.

[bondhugula]

LGTM

Nit: ... also be made to be MemRefsNormalizable -> should have the MemRefsNormalizable trait.

mlir/test/Transforms/normalize-memrefs.mlir

327

Nit: no need of the {{.*}} = LHS part.

[adstraw]

Thank you @bondhugula.

[bondhugula]

LGTM.

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

let description = [{
     ```mlir
     %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
     ```
   }];
 }
 
 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
     parallel band, defining multiple SSA values for its induction variables. It
     has one region capturing the parallel band body. The induction variables are
     represented as arguments of this region. These SSA values always have type
     index, which is the size of the machine word. The strides, represented by
     steps, are positive constant integers which defaults to "1" if not present.
def AffineStoreOp : AffineStoreOpBase<"store"> {
   ];
 
   let extraClassDeclaration = extraClassDeclarationBase;
 
   let hasCanonicalizer = 1;
   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
     terminates the region. The semantics of how the values yielded are used
     is defined by the parent operation.
     If "affine.yield" has any operands, the operands must match the parent
     operation's results.
     If the parent operation defines no values, then the "affine.yield" may be

mlir/lib/Transforms/NormalizeMemRefs.cpp

for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
     resultTypes.push_back(newMemRefType);
     resultTypeNormalized = true;
   }
   result.addTypes(resultTypes);
   // When all of the results of `oldOp` have no memrefs or memrefs without
   // 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

 // CHECK-LABEL: func @use_value_of_external
 // CHECK-SAME: (%[[A:arg[0-9]+]]: memref<4x4xf64>, %[[B:arg[0-9]+]]: f64) -> memref<2x4xf64>
 func @use_value_of_external(%A: memref<16xf64, #tile>, %B: f64) -> (memref<8xf64, #tile>) {
   %res = call @external_func_B(%A, %B) : (memref<16xf64, #tile>, f64) -> (memref<8xf64, #tile>)
   return %res : memref<8xf64, #tile>
 }
 // 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>
+}