[flang-commits] [flang] 5156d1a - [flang][hlfir] get constant extents when possible

[Tom Eccles via flang-commits]

Author: Tom Eccles
Date: 2023-04-25T12:43:48Z
New Revision: 5156d1a797be2c4a6af40fb33a574c3d379196ae

URL: https://github.com/llvm/llvm-project/commit/5156d1a797be2c4a6af40fb33a574c3d379196ae
DIFF: https://github.com/llvm/llvm-project/commit/5156d1a797be2c4a6af40fb33a574c3d379196ae.diff

LOG: [flang][hlfir] get constant extents when possible

If we know the extent at compile time, it is better to generate an
arith.constant than hlfir.get_extent. This gives more information earlier
in compilation (before hflir.get_extent is lowered).

Reviewed By: jeanPerier

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

Added: 
    

Modified: 
    flang/lib/Optimizer/Builder/HLFIRTools.cpp
    flang/test/HLFIR/extents-of-shape-of.f90

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Builder/HLFIRTools.cpp b/flang/lib/Optimizer/Builder/HLFIRTools.cpp
index 9b9174eea80ae..f4481734e9a50 100644
--- a/flang/lib/Optimizer/Builder/HLFIRTools.cpp
+++ b/flang/lib/Optimizer/Builder/HLFIRTools.cpp
@@ -36,11 +36,22 @@ getExplicitExtentsFromShape(mlir::Value shape, fir::FirOpBuilder &builder) {
   } else if (mlir::dyn_cast_or_null<fir::ShiftOp>(shapeOp)) {
     return {};
   } else if (auto s = mlir::dyn_cast_or_null<hlfir::ShapeOfOp>(shapeOp)) {
+    hlfir::ExprType expr = s.getExpr().getType().cast<hlfir::ExprType>();
+    llvm::ArrayRef<int64_t> exprShape = expr.getShape();
+    mlir::Type indexTy = builder.getIndexType();
     fir::ShapeType shapeTy = shape.getType().cast<fir::ShapeType>();
     result.reserve(shapeTy.getRank());
     for (unsigned i = 0; i < shapeTy.getRank(); ++i) {
-      auto op = builder.create<hlfir::GetExtentOp>(shape.getLoc(), shape, i);
-      result.emplace_back(op.getResult());
+      int64_t extent = exprShape[i];
+      mlir::Value extentVal;
+      if (extent == expr.getUnknownExtent()) {
+        auto op = builder.create<hlfir::GetExtentOp>(shape.getLoc(), shape, i);
+        extentVal = op.getResult();
+      } else {
+        extentVal =
+            builder.createIntegerConstant(shape.getLoc(), indexTy, extent);
+      }
+      result.emplace_back(extentVal);
     }
   } else {
     TODO(shape.getLoc(), "read fir.shape to get extents");

diff  --git a/flang/test/HLFIR/extents-of-shape-of.f90 b/flang/test/HLFIR/extents-of-shape-of.f90
index 85479ebb407f6..8fc7133985ffc 100644
--- a/flang/test/HLFIR/extents-of-shape-of.f90
+++ b/flang/test/HLFIR/extents-of-shape-of.f90
@@ -1,7 +1,7 @@
 ! RUN: bbc -emit-fir -hlfir %s -o - | FileCheck --check-prefix CHECK-ALL --check-prefix CHECK-HLFIR %s
 ! RUN: bbc -emit-fir -hlfir %s -o - | fir-opt --lower-hlfir-intrinsics | fir-opt --bufferize-hlfir | fir-opt --convert-hlfir-to-fir | FileCheck --check-prefix CHECK-ALL --check-prefix CHECK-FIR %s
 subroutine foo(a, b)
-  real :: a(:, :), b(:, :)
+  real :: a(2, 2), b(:, :)
   interface
     elemental subroutine elem_sub(x)
       real, intent(in) :: x
@@ -10,19 +10,19 @@ elemental subroutine elem_sub(x)
   call elem_sub(matmul(a, b))
 end subroutine
 ! CHECK-ALL-LABEL:   func.func @_QPfoo
-! CHECK-ALL:             %[[A_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {fir.bindc_name = "a"}
+! CHECK-ALL:             %[[A_ARG:.*]]: !fir.ref<!fir.array<2x2xf32>> {fir.bindc_name = "a"}
 ! CHECK-ALL:             %[[B_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {fir.bindc_name = "b"}
 
 ! CHECK-HLFIR-DAG:     %[[A_VAR:.*]]:2 = hlfir.declare %[[A_ARG]]
 ! CHECK-HLFIR-DAG:     %[[B_VAR:.*]]:2 = hlfir.declare %[[B_ARG]]
 ! CHECK-HLFIR-NEXT:    %[[MUL:.*]] = hlfir.matmul %[[A_VAR]]#0 %[[B_VAR]]#0
-! CHECK-HLFIR-NEXT:    %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr<?x?xf32>) -> !fir.shape<2>
-! CHECK-HLFIR-NEXT:    %[[EXT0:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 0 : index} : (!fir.shape<2>) -> index
+! CHECK-HLFIR-NEXT:    %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr<2x?xf32>) -> !fir.shape<2>
+! CHECK-HLFIR-NEXT:    %[[EXT0:.*]] = arith.constant 2 : index
 ! CHECK-HLFIR-NEXT:    %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<2>) -> index
 ! CHECK-HLFIR-NEXT:    %[[C1:.*]] = arith.constant 1 : index
 ! CHECK-HLFIR-NEXT:    fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] {
 ! CHECK-HLFIR-NEXT:      fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] {
-! CHECK-HLFIR-NEXT:        %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr<?x?xf32>, index, index) -> f32
+! CHECK-HLFIR-NEXT:        %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr<2x?xf32>, index, index) -> f32
 ! CHECK-HLFIR-NEXT:        %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
 ! CHECK-HLFIR-NEXT:        fir.call
 ! CHECK-HLFIR-NEXT:        hlfir.end_associate
@@ -39,9 +39,10 @@ elemental subroutine elem_sub(x)
 ! CHECK-FIR-NEXT:      %[[DIMS1:.*]]:3 = fir.box_dims %[[MUL]], %[[C1]]
 ! ...
 ! CHECK-FIR:           %[[SHAPE:.*]] = fir.shape %[[DIMS0]]#1, %[[DIMS1]]#1
+! CHECK-FIR-NEXT:      %[[C2:.*]] = arith.constant 2 : index
 ! CHECK-FIR-NEXT:      %[[C1_1:.*]] = arith.constant 1 : index
 ! CHECK-FIR-NEXT:      fir.do_loop %[[ARG2:.*]] = %[[C1_1]] to %[[DIMS1]]#1 step %[[C1_1]] {
-! CHECK-FIR-NEXT:        fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[DIMS0]]#1 step %[[C1_1]] {
+! CHECK-FIR-NEXT:        fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[C2]] step %[[C1_1]] {
 ! ...
 
 ! CHECK-ALL:           return