[flang-commits] [flang] 775de67 - [flang] convert stack arrays allocation to match old type

[Tom Eccles via flang-commits]

Author: Tom Eccles
Date: 2023-06-05T09:57:57Z
New Revision: 775de6754af8d5122f10204097f41be02c5c9f3b

URL: https://github.com/llvm/llvm-project/commit/775de6754af8d5122f10204097f41be02c5c9f3b
DIFF: https://github.com/llvm/llvm-project/commit/775de6754af8d5122f10204097f41be02c5c9f3b.diff

LOG: [flang] convert stack arrays allocation to match old type

The old fir.allocmem operation returned a !fir.heap<.> type. The new
fir.alloca operation returns a !fir.ref<.> type. This patch inserts a
fir.convert so that the old type is preserved. This prevents verifier
failures when types returned from fir.if statements don't match the
expected type.

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

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/StackArrays.cpp
    flang/test/Transforms/stack-arrays.fir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index 0f21e755dad27..f01c94d977a1a 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -465,6 +465,29 @@ StackArraysAnalysisWrapper::getCandidateOps(mlir::Operation *func) {
   return &funcMaps[func];
 }
 
+/// Restore the old allocation type exected by existing code
+static mlir::Value convertAllocationType(mlir::PatternRewriter &rewriter,
+                                         const mlir::Location &loc,
+                                         mlir::Value heap, mlir::Value stack) {
+  mlir::Type heapTy = heap.getType();
+  mlir::Type stackTy = stack.getType();
+
+  if (heapTy == stackTy)
+    return stack;
+
+  fir::HeapType firHeapTy = mlir::cast<fir::HeapType>(heapTy);
+  fir::ReferenceType firRefTy = mlir::cast<fir::ReferenceType>(stackTy);
+  assert(firHeapTy.getElementType() == firRefTy.getElementType() &&
+         "Allocations must have the same type");
+
+  auto insertionPoint = rewriter.saveInsertionPoint();
+  rewriter.setInsertionPointAfter(stack.getDefiningOp());
+  mlir::Value conv =
+      rewriter.create<fir::ConvertOp>(loc, firHeapTy, stack).getResult();
+  rewriter.restoreInsertionPoint(insertionPoint);
+  return conv;
+}
+
 mlir::LogicalResult
 AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
                                     mlir::PatternRewriter &rewriter) const {
@@ -485,7 +508,9 @@ AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
     rewriter.eraseOp(erase);
 
   // replace references to heap allocation with references to stack allocation
-  rewriter.replaceAllUsesWith(allocmem.getResult(), alloca->getResult());
+  mlir::Value newValue = convertAllocationType(
+      rewriter, allocmem.getLoc(), allocmem.getResult(), alloca->getResult());
+  rewriter.replaceAllUsesWith(allocmem.getResult(), newValue);
 
   // remove allocmem operation
   rewriter.eraseOp(allocmem.getOperation());

diff  --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index 046a402831aa8..eaea7009e9f30 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -103,10 +103,13 @@ func.func @dfa3a(%arg0: i1) {
 }
 // CHECK:     func.func @dfa3a(%arg0: i1) {
 // CHECK-NEXT:  %[[MEM:.*]] = fir.alloca !fir.array<1xi8>
+// CHECK-NEXT:  %[[HEAP:.*]] = fir.convert %[[MEM]] : (!fir.ref<!fir.array<1xi8>>) -> !fir.heap<!fir.array<1xi8>>
 // CHECK-NEXT:  fir.if %arg0 {
-// CHECK-NEXT:    func.call @dfa3a_foo(%[[MEM]])
+// CHECK-NEXT:    %[[REF:.*]] = fir.convert %[[HEAP]] : (!fir.heap<!fir.array<1xi8>>) -> !fir.ref<!fir.array<1xi8>>
+// CHECK-NEXT:    func.call @dfa3a_foo(%[[REF]])
 // CHECK-NEXT:  } else {
-// CHECK-NEXT:    func.call @dfa3a_bar(%[[MEM]])
+// CHECK-NEXT:    %[[REF:.*]] = fir.convert %[[HEAP]] : (!fir.heap<!fir.array<1xi8>>) -> !fir.ref<!fir.array<1xi8>>
+// CHECK-NEXT:    func.call @dfa3a_bar(%[[REF]])
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 // CHECK-NEXT:  }