[flang-commits] [flang] [flang][fir] add fir.is_assumed_size codegen (PR #93953)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/93953

Lower to IsAssumedSize runtime call in `--fir-assumed-rank-op` pass.

>From 1e9ecebeec3e8ab62f09f4ac99eea209c874da26 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Fri, 31 May 2024 05:11:34 -0700
Subject: [PATCH] [flang][fir] add fir.is_assumed_size codegen

---
 .../flang/Optimizer/Builder/Runtime/Support.h |  4 +++
 .../lib/Optimizer/Builder/Runtime/Support.cpp | 10 +++++++
 .../Transforms/AssumedRankOpConversion.cpp    | 26 +++++++++++++++++++
 flang/test/Fir/is_assumed_size-codegen.fir    | 25 ++++++++++++++++++
 4 files changed, 65 insertions(+)
 create mode 100644 flang/test/Fir/is_assumed_size-codegen.fir

diff --git a/flang/include/flang/Optimizer/Builder/Runtime/Support.h b/flang/include/flang/Optimizer/Builder/Runtime/Support.h
index fe263ca2975ee..d0a474d75d2eb 100644
--- a/flang/include/flang/Optimizer/Builder/Runtime/Support.h
+++ b/flang/include/flang/Optimizer/Builder/Runtime/Support.h
@@ -27,5 +27,9 @@ void genCopyAndUpdateDescriptor(fir::FirOpBuilder &builder, mlir::Location loc,
                                 mlir::Value newAttribute,
                                 mlir::Value newLowerBounds);
 
+/// Generate call to `IsAssumedSize` runtime routine.
+mlir::Value genIsAssumedSize(fir::FirOpBuilder &builder, mlir::Location loc,
+                             mlir::Value box);
+
 } // namespace fir::runtime
 #endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_SUPPORT_H
diff --git a/flang/lib/Optimizer/Builder/Runtime/Support.cpp b/flang/lib/Optimizer/Builder/Runtime/Support.cpp
index 12e47233e3d99..b5e9ddb87c7c4 100644
--- a/flang/lib/Optimizer/Builder/Runtime/Support.cpp
+++ b/flang/lib/Optimizer/Builder/Runtime/Support.cpp
@@ -44,3 +44,13 @@ void fir::runtime::genCopyAndUpdateDescriptor(fir::FirOpBuilder &builder,
   }
   builder.create<fir::CallOp>(loc, func, args);
 }
+
+mlir::Value fir::runtime::genIsAssumedSize(fir::FirOpBuilder &builder,
+                                           mlir::Location loc,
+                                           mlir::Value box) {
+  mlir::func::FuncOp func =
+      fir::runtime::getRuntimeFunc<mkRTKey(IsAssumedSize)>(loc, builder);
+  auto fTy = func.getFunctionType();
+  auto args = fir::runtime::createArguments(builder, loc, fTy, box);
+  return builder.create<fir::CallOp>(loc, func, args).getResult(0);
+}
diff --git a/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp b/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp
index 5cc70c4d61257..2c545d66ebd8e 100644
--- a/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/AssumedRankOpConversion.cpp
@@ -112,6 +112,31 @@ class ReboxAssumedRankConv
   fir::KindMapping kindMap;
 };
 
+class IsAssumedSizeConv : public mlir::OpRewritePattern<fir::IsAssumedSizeOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  IsAssumedSizeConv(mlir::MLIRContext *context, mlir::SymbolTable *symbolTable,
+                    fir::KindMapping kindMap)
+      : mlir::OpRewritePattern<fir::IsAssumedSizeOp>(context),
+        symbolTable{symbolTable}, kindMap{kindMap} {};
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::IsAssumedSizeOp isAssumedSizeOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    fir::FirOpBuilder builder{rewriter, kindMap, symbolTable};
+    mlir::Location loc = isAssumedSizeOp.getLoc();
+    mlir::Value result =
+        fir::runtime::genIsAssumedSize(builder, loc, isAssumedSizeOp.getVal());
+    rewriter.replaceOp(isAssumedSizeOp, result);
+    return mlir::success();
+  }
+
+private:
+  mlir::SymbolTable *symbolTable = nullptr;
+  fir::KindMapping kindMap;
+};
+
 /// Convert FIR structured control flow ops to CFG ops.
 class AssumedRankOpConversion
     : public fir::impl::AssumedRankOpConversionBase<AssumedRankOpConversion> {
@@ -123,6 +148,7 @@ class AssumedRankOpConversion
     fir::KindMapping kindMap = fir::getKindMapping(mod);
     mlir::RewritePatternSet patterns(context);
     patterns.insert<ReboxAssumedRankConv>(context, &symbolTable, kindMap);
+    patterns.insert<IsAssumedSizeConv>(context, &symbolTable, kindMap);
     mlir::GreedyRewriteConfig config;
     config.enableRegionSimplification = false;
     (void)applyPatternsAndFoldGreedily(mod, std::move(patterns), config);
diff --git a/flang/test/Fir/is_assumed_size-codegen.fir b/flang/test/Fir/is_assumed_size-codegen.fir
new file mode 100644
index 0000000000000..646446c00e879
--- /dev/null
+++ b/flang/test/Fir/is_assumed_size-codegen.fir
@@ -0,0 +1,25 @@
+// Test fir.is_assumed_size lowering to runtime calls in fir-assumed-rank-op pass.
+// RUN: fir-opt -o - --fir-assumed-rank-op %s | FileCheck %s
+
+func.func @test_ranked(%arg0: !fir.class<!fir.array<?xnone>> ) -> i1 {
+  %1 = fir.is_assumed_size %arg0 : (!fir.class<!fir.array<?xnone>>) -> i1
+  return %1 : i1
+}
+
+func.func @test_assumed_rank(%arg0: !fir.box<!fir.array<*:f32>> ) -> i1 {
+  %1 = fir.is_assumed_size %arg0 : (!fir.box<!fir.array<*:f32>>) -> i1
+  return %1 : i1
+}
+// CHECK-LABEL:   func.func @test_ranked(
+// CHECK-SAME:                           %[[VAL_0:.*]]: !fir.class<!fir.array<?xnone>>) -> i1 {
+// CHECK:           %[[VAL_1:.*]] = fir.convert %[[VAL_0]] : (!fir.class<!fir.array<?xnone>>) -> !fir.box<none>
+// CHECK:           %[[VAL_2:.*]] = fir.call @_FortranAIsAssumedSize(%[[VAL_1]]) : (!fir.box<none>) -> i1
+// CHECK:           return %[[VAL_2]] : i1
+// CHECK:         }
+
+// CHECK-LABEL:   func.func @test_assumed_rank(
+// CHECK-SAME:                                 %[[VAL_0:.*]]: !fir.box<!fir.array<*:f32>>) -> i1 {
+// CHECK:           %[[VAL_1:.*]] = fir.convert %[[VAL_0]] : (!fir.box<!fir.array<*:f32>>) -> !fir.box<none>
+// CHECK:           %[[VAL_2:.*]] = fir.call @_FortranAIsAssumedSize(%[[VAL_1]]) : (!fir.box<none>) -> i1
+// CHECK:           return %[[VAL_2]] : i1
+// CHECK:         }