[FlattenMemRef] Flatten MemRef AllocaOp (#8352)

Author: KelvinChung2000

lib/Transforms/FlattenMemRefs.cpp

@@ -178,6 +178,21 @@ struct AllocOpConversion : public OpConversionPattern<memref::AllocOp> {
   }
 };
 
+struct AllocaOpConversion : public OpConversionPattern<memref::AllocaOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::AllocaOp op, OpAdaptor /*adaptor*/,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType type = op.getType();
+    if (isUniDimensional(type) || !type.hasStaticShape())
+      return failure();
+    MemRefType newType = getFlattenedMemRefType(type);
+    rewriter.replaceOpWithNewOp<memref::AllocaOp>(op, newType);
+    return success();
+  }
+};
+
 struct GlobalOpConversion : public OpConversionPattern<memref::GlobalOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -352,6 +367,8 @@ static void populateFlattenMemRefsLegality(ConversionTarget &target) {
   target.addLegalDialect<arith::ArithDialect>();
   target.addDynamicallyLegalOp<memref::AllocOp>(
       [](memref::AllocOp op) { return isUniDimensional(op.getType()); });
+  target.addDynamicallyLegalOp<memref::AllocaOp>(
+      [](memref::AllocaOp op) { return isUniDimensional(op.getType()); });
   target.addDynamicallyLegalOp<memref::StoreOp>(
       [](memref::StoreOp op) { return op.getIndices().size() == 1; });
   target.addDynamicallyLegalOp<memref::LoadOp>(
@@ -426,8 +443,8 @@ struct FlattenMemRefPass
     RewritePatternSet patterns(ctx);
     SetVector<StringRef> rewrittenCallees;
     patterns.add<LoadOpConversion, StoreOpConversion, AllocOpConversion,
-                 GlobalOpConversion, GetGlobalOpConversion, ReshapeOpConversion,
-                 OperandConversionPattern<func::ReturnOp>,
+                 AllocaOpConversion, GlobalOpConversion, GetGlobalOpConversion,
+                 ReshapeOpConversion, OperandConversionPattern<func::ReturnOp>,
                  OperandConversionPattern<memref::DeallocOp>,
                  CondBranchOpConversion,
                  OperandConversionPattern<memref::DeallocOp>,

test/Transforms/flatten_memref.mlir

@@ -314,3 +314,14 @@ module {
     return
   }
 }
+
+// -----
+ 
+ // CHECK-LABEL:   func @allocas() -> memref<16xi32> {
+ // CHECK:           %[[VAL_0:.*]] = memref.alloca() : memref<16xi32>
+ // CHECK:           return %[[VAL_0]] : memref<16xi32>
+ // CHECK:         }
+ func.func @allocas() -> memref<4x4xi32> {
+   %0 = memref.alloca() : memref<4x4xi32>
+   return %0 : memref<4x4xi32>
+ }