From b8708753c89db07eb6993559b4546e5a97d94530 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin@nvidia.com>
Date: Wed, 29 Jan 2025 12:01:13 -0800
Subject: [PATCH] [flang] Allow non-index length parameter on exprs fed into
 hlfir.get_length. (#124827)

The length might be any integer, so hlfir.get_length lowering
should explicitly cast it to `index`.
---
 .../HLFIR/Transforms/BufferizeHLFIR.cpp       |  1 +
 flang/test/HLFIR/get_length_codegen.fir       | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
index 48e3db0eb39b..e664834d31d3 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -362,6 +362,7 @@ struct GetLengthOpConversion
     if (!length)
       return rewriter.notifyMatchFailure(
           getLength, "could not deduce length from GetLengthOp operand");
+    length = builder.createConvert(loc, builder.getIndexType(), length);
     rewriter.replaceOp(getLength, length);
     return mlir::success();
   }
diff --git a/flang/test/HLFIR/get_length_codegen.fir b/flang/test/HLFIR/get_length_codegen.fir
index 7cb36d2cd3c6..c828aa5653d6 100644
--- a/flang/test/HLFIR/get_length_codegen.fir
+++ b/flang/test/HLFIR/get_length_codegen.fir
@@ -30,3 +30,22 @@ fir.global linkonce @_QQclX616263 constant : !fir.char<1,3> {
 // CHECK:           %[[VAL_33:.*]]:2 = hlfir.declare %[[VAL_31:.*]] typeparams %[[VAL_9]] {uniq_name = ".tmp"} : (!fir.ref<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
 // CHECK:           return %[[VAL_9]] : index
 // CHECK:         }
+
+// Test get_length taking the length from an expression with i32 length parameter.
+func.func @i32_length(%char: !fir.boxchar<1>, %shape : i32, %len : i32) -> index {
+  %14 = fir.shape %shape : (i32) -> !fir.shape<1>
+  %15 = hlfir.elemental %14 typeparams %len unordered : (!fir.shape<1>, i32) -> !hlfir.expr<?x!fir.char<1,?>> {
+  ^bb0(%arg0: index):
+    hlfir.yield_element %char : !fir.boxchar<1>
+  }
+  %18 = hlfir.get_length %15 : (!hlfir.expr<?x!fir.char<1,?>>) -> index
+  hlfir.destroy %15 : !hlfir.expr<?x!fir.char<1,?>>
+  return %18 : index
+}
+// CHECK-LABEL:   func.func @i32_length(
+// CHECK-SAME:                          %[[VAL_0:.*]]: !fir.boxchar<1>,
+// CHECK-SAME:                          %[[VAL_1:.*]]: i32,
+// CHECK-SAME:                          %[[VAL_2:.*]]: i32) -> index {
+// CHECK:           %[[VAL_14:.*]] = fir.convert %[[VAL_2]] : (i32) -> index
+// CHECK:           return %[[VAL_14]] : index
+// CHECK:         }