[flang] Support non-index shape in hlfir.get_extent codegen.

[vzakhari]

hlfir.reshape inlining uncovered an existing bug that non-index shapes
result in failures during hlfir.get_extent conversion to FIR.
I could have "fixed" the shape during hlfir.reshape inlining,
but this seems like a better fix.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Slava Zakharin (vzakhari)

Changes

hlfir.reshape inlining uncovered an existing bug that non-index shapes
result in failures during hlfir.get_extent conversion to FIR.
I could have "fixed" the shape during hlfir.reshape inlining,
but this seems like a better fix.

---

Full diff: https://github.com/llvm/llvm-project/pull/124622.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp (+3)
• (added) flang/test/HLFIR/getextent-codegen.fir (+22)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
index 536f2077e4f700..496a5560ac615c 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
@@ -735,6 +735,9 @@ class GetExtentOpConversion
       llvm::APInt dim = getExtentOp.getDim();
       uint64_t dimVal = dim.getLimitedValue(shapeTy.getRank());
       mlir::Value extent = s.getExtents()[dimVal];
+      fir::FirOpBuilder builder(rewriter, getExtentOp.getOperation());
+      extent = builder.createConvert(getExtentOp.getLoc(),
+                                     builder.getIndexType(), extent);
       rewriter.replaceOp(getExtentOp, extent);
       return mlir::success();
     }
diff --git a/flang/test/HLFIR/getextent-codegen.fir b/flang/test/HLFIR/getextent-codegen.fir
new file mode 100644
index 00000000000000..4f48569424e1b4
--- /dev/null
+++ b/flang/test/HLFIR/getextent-codegen.fir
@@ -0,0 +1,22 @@
+// Test hlfir.get_extent code generation to FIR
+// RUN: fir-opt %s -convert-hlfir-to-fir | FileCheck %s
+
+func.func @test1(%arg0: i1, %arg1: i32, %arg2: i64, %arg3: index) -> (index, index, index, index) {
+  %0 = fir.shape %arg0, %arg1, %arg2, %arg3 : (i1, i32, i64, index) -> !fir.shape<4>
+  %1 = hlfir.get_extent %0 {dim = 0 : index} : (!fir.shape<4>) -> index
+  %2 = hlfir.get_extent %0 {dim = 1 : index} : (!fir.shape<4>) -> index
+  %3 = hlfir.get_extent %0 {dim = 2 : index} : (!fir.shape<4>) -> index
+  %4 = hlfir.get_extent %0 {dim = 3 : index} : (!fir.shape<4>) -> index
+  return %1, %2, %3, %4 : index, index, index, index
+}
+// CHECK-LABEL:   func.func @test1(
+// CHECK-SAME:                     %[[VAL_0:.*]]: i1,
+// CHECK-SAME:                     %[[VAL_1:.*]]: i32,
+// CHECK-SAME:                     %[[VAL_2:.*]]: i64,
+// CHECK-SAME:                     %[[VAL_3:.*]]: index) -> (index, index, index, index) {
+// CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] : (i1, i32, i64, index) -> !fir.shape<4>
+// CHECK:           %[[VAL_5:.*]] = fir.convert %[[VAL_0]] : (i1) -> index
+// CHECK:           %[[VAL_6:.*]] = fir.convert %[[VAL_1]] : (i32) -> index
+// CHECK:           %[[VAL_7:.*]] = fir.convert %[[VAL_2]] : (i64) -> index
+// CHECK:           return %[[VAL_5]], %[[VAL_6]], %[[VAL_7]], %[[VAL_3]] : index, index, index, index
+// CHECK:         }

[jeanPerier]

Thank you @vzakhari!

[tblah]

LGTM, thanks!