From 5228c2cbd6d834a76c24ca9f2aaa2c4f7a0b5510 Mon Sep 17 00:00:00 2001
From: jeanPerier <jperier@nvidia.com>
Date: Fri, 31 May 2024 08:38:40 +0200
Subject: [PATCH] [flang][FIR] add fir.is_assumed_size operation (#93853)

Assumed-rank fir.box/class may describe assumed-size array. This case
needs special handling in SELECT RANK. It is not possible to generate
FIR code to detect that a fir.box is an assumed-size (the way to detect
that is to check that upper dimension extent is -1 in the descriptor).

Instead of emitting a runtime call directly in lowering, add an
operation that can later be lowered to a runtime call or inline code
when the descriptor layout is known.
---
 flang/include/flang/Optimizer/Dialect/FIROps.td | 17 +++++++++++++++++
 flang/test/Fir/fir-ops.fir                      | 11 +++++++++++
 flang/test/Fir/invalid.fir                      |  8 ++++++++
 3 files changed, 36 insertions(+)

diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 3afc97475db11b..3b9512e07b035f 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -1171,6 +1171,23 @@ def fir_BoxIsArrayOp : fir_SimpleOp<"box_isarray", [NoMemoryEffect]> {
   let results = (outs BoolLike);
 }
 
+def fir_IsAssumedSizeOp : fir_SimpleOp<"is_assumed_size", [NoMemoryEffect]> {
+  let summary = "detect if a boxed value is an assumed-size array";
+
+  let description = [{
+    Fir box SSA values may describe assumed-size arrays. This operation
+    allows detecting this, even for assumed-rank box.
+
+    ```
+      %a = fir.is_assumed_size %b : (!fir.box<!fir.array<*:f64>>) -> i1
+    ```
+  }];
+
+  let arguments = (ins BoxOrClassType:$val);
+
+  let results = (outs BoolLike);
+}
+
 def fir_BoxIsPtrOp : fir_SimpleOp<"box_isptr", [NoMemoryEffect]> {
   let summary = "is the boxed value a POINTER?";
 
diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir
index a826dd49ef99db..ecc5c4ecff35df 100644
--- a/flang/test/Fir/fir-ops.fir
+++ b/flang/test/Fir/fir-ops.fir
@@ -912,3 +912,14 @@ func.func @test_rebox_assumed_rank(%arg0: !fir.box<!fir.array<*:f32>> ) {
   // CHECK: fir.rebox_assumed_rank %[[A]] lbs ones : (!fir.box<!fir.array<*:f32>>) -> !fir.box<!fir.array<*:f32>>
   // CHECK: fir.rebox_assumed_rank %[[A]] lbs zeroes : (!fir.box<!fir.array<*:f32>>) -> !fir.box<!fir.array<*:f32>>
   // CHECK: fir.rebox_assumed_rank %[[A]] lbs preserve : (!fir.box<!fir.array<*:f32>>) -> !fir.box<!fir.array<*:f32>>
+
+func.func @test_is_assumed_size(%arg0: !fir.class<!fir.array<*:none>>, %arg1 : !fir.box<!fir.array<?xf32>>) {
+  %1 = fir.is_assumed_size %arg0 : (!fir.class<!fir.array<*:none>>) -> i1
+  %2 = fir.is_assumed_size %arg1 : (!fir.box<!fir.array<?xf32>>) -> i1
+  return
+}
+// CHECK-LABEL: func.func @test_is_assumed_size(
+// CHECK-SAME: %[[A:.*]]: !fir.class<!fir.array<*:none>>,
+// CHECK-SAME: %[[B:.*]]: !fir.box<!fir.array<?xf32>>)
+  // CHECK: fir.is_assumed_size %[[A]] : (!fir.class<!fir.array<*:none>>) -> i1
+  // CHECK: fir.is_assumed_size %[[B]] : (!fir.box<!fir.array<?xf32>>) -> i1
diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir
index f1e1aa433b9b03..086a426db5642e 100644
--- a/flang/test/Fir/invalid.fir
+++ b/flang/test/Fir/invalid.fir
@@ -1002,3 +1002,11 @@ func.func @bad_rebox_assumed_rank_3(%arg0: !fir.box<!fir.array<*:f32>> ) {
   %1 = fir.rebox_assumed_rank %arg0 lbs ones : (!fir.box<!fir.array<*:f32>>) -> !fir.box<!fir.array<*:i32>>
   return
 }
+
+// -----
+
+func.func @bad_is_assumed_size(%arg0: !fir.ref<!fir.array<*:none>>) {
+  // expected-error@+1{{op operand #0 must be box or class, but got '!fir.ref<!fir.array<*:none>>'}}
+  %1 = fir.is_assumed_size %arg0 : (!fir.ref<!fir.array<*:none>>) -> i1
+  return
+}