[flang][hlfir] get extents from hlfir.shape_of

[tblah]

Summary

If the extents were known, this should have been canonicalised into a fir.shape operation. Therefore, the extents at this point are not known at compile time. Use hlfir.get_extents to delay resolving the real extent until after the expression is bufferized.

Depends on: D146831 D148220

[jeanPerier]

The situation is different between a fir.box and an hlfir.expr here. Returning an empty vector from a fir.shape value here will violate what some of the user expect (like getIndexExtents). I think a new operation is needed here to get an extent from a fir.shape instead (or maybe from an hlfir.expr).
The rational is that it must always be possible to get an extent value in lowering (required to generate loop nest, or answer extent inquires). With a fir.box, if this code "fails" because the input is a fir::ShiftOp, the helper tools will generate fir.box_dim on the fir.box. But with an hlfir.expr, there is no way to get a single extent.

Places that would fail:

This will violates getIndexExtents contract: https://github.com/llvm/llvm-project/blob/main/flang/include/flang/Optimizer/Builder/HLFIRTools.h#L306-L310
Example of program that would likely fail to compile properly:

subroutine foo(a, b)
 real :: a(:, :), b(:, :)
 interface
   elemental subroutine elem_sub(x)
    real, intent(in) :: x
   end subroutine
 end interface
 call elem_sub(matmul(a,b))
end subroutine

This would also still leave the TODO in genExtent: https://github.com/llvm/llvm-project/blob/3ad26060e4bceb2cf9f4959d659cbb29d88344cf/flang/lib/Optimizer/Builder/HLFIRTools.cpp#L523
Likely hit by:

subroutine foo(a, b)
 real :: a(:, :), b(:, :)
 print *, size(matmul(a,b), dim=1)
end subroutine

[tblah]

Updating to get extents using a new hlfir.get_extent operation

[tblah]

In D146832#4223300, @jeanPerier wrote:

The situation is different between a fir.box and an hlfir.expr here. Returning an empty vector from a fir.shape value here will violate what some of the user expect (like getIndexExtents). I think a new operation is needed here to get an extent from a fir.shape instead (or maybe from an hlfir.expr).
The rational is that it must always be possible to get an extent value in lowering (required to generate loop nest, or answer extent inquires). With a fir.box, if this code "fails" because the input is a fir::ShiftOp, the helper tools will generate fir.box_dim on the fir.box. But with an hlfir.expr, there is no way to get a single extent.

Places that would fail:

This will violates getIndexExtents contract: https://github.com/llvm/llvm-project/blob/main/flang/include/flang/Optimizer/Builder/HLFIRTools.h#L306-L310
Example of program that would likely fail to compile properly:

subroutine foo(a, b)
 real :: a(:, :), b(:, :)
 interface
   elemental subroutine elem_sub(x)
    real, intent(in) :: x
   end subroutine
 end interface
 call elem_sub(matmul(a,b))
end subroutine

This would also still leave the TODO in genExtent: https://github.com/llvm/llvm-project/blob/3ad26060e4bceb2cf9f4959d659cbb29d88344cf/flang/lib/Optimizer/Builder/HLFIRTools.cpp#L523
Likely hit by:

subroutine foo(a, b)
 real :: a(:, :), b(:, :)
 print *, size(matmul(a,b), dim=1)
end subroutine

Thanks for pointing these out. I added the first example as a testcase here and in https://reviews.llvm.org/D148222

[jeanPerier]

There seems to be something wrong with the patch application. Other than that, this looks great, thanks for dealing with the issue.

[tblah]

Bumping without change to re-run CI

flang/lib/Optimizer/Builder/HLFIRTools.cpp

 #include "mlir/IR/IRMapping.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include <optional>
 
 // Return explicit extents. If the base is a fir.box, this won't read it to
 // return the extents and will instead return an empty vector.
 static llvm::SmallVector<mlir::Value>
-getExplicitExtentsFromShape(mlir::Value shape) {
+getExplicitExtentsFromShape(mlir::Value shape, fir::FirOpBuilder &builder) {
   llvm::SmallVector<mlir::Value> result;
   auto *shapeOp = shape.getDefiningOp();
   if (auto s = mlir::dyn_cast_or_null<fir::ShapeOp>(shapeOp)) {
     auto e = s.getExtents();
     result.append(e.begin(), e.end());
   } else if (auto s = mlir::dyn_cast_or_null<fir::ShapeShiftOp>(shapeOp)) {
     auto e = s.getExtents();
     result.append(e.begin(), e.end());
   } else if (mlir::dyn_cast_or_null<fir::ShiftOp>(shapeOp)) {
     return {};
+  } else if (auto s = mlir::dyn_cast_or_null<hlfir::ShapeOfOp>(shapeOp)) {
+    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());
+    }
   } else {
     TODO(shape.getLoc(), "read fir.shape to get extents");
   }
   return result;
 }
 static llvm::SmallVector<mlir::Value>
-getExplicitExtents(fir::FortranVariableOpInterface var) {
+getExplicitExtents(fir::FortranVariableOpInterface var,
+                   fir::FirOpBuilder &builder) {
   if (mlir::Value shape = var.getShape())
-    return getExplicitExtentsFromShape(var.getShape());
+    return getExplicitExtentsFromShape(var.getShape(), builder);
   return {};
 }
 
 // Return explicit lower bounds. For pointers and allocatables, this will not
 // read the lower bounds and instead return an empty vector.
 static llvm::SmallVector<mlir::Value>
 getExplicitLboundsFromShape(mlir::Value shape) {
   llvm::SmallVector<mlir::Value> result;

 }
 
 llvm::SmallVector<std::pair<mlir::Value, mlir::Value>>
 hlfir::genBounds(mlir::Location loc, fir::FirOpBuilder &builder,
                  mlir::Value shape) {
   assert((shape.getType().isa<fir::ShapeShiftType>() ||
           shape.getType().isa<fir::ShapeType>()) &&
          "shape must contain extents");
-  auto extents = getExplicitExtentsFromShape(shape);
+  auto extents = getExplicitExtentsFromShape(shape, builder);
   auto lowers = getExplicitLboundsFromShape(shape);
   assert(lowers.empty() || lowers.size() == extents.size());
   mlir::Type idxTy = builder.getIndexType();
   mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
   llvm::SmallVector<std::pair<mlir::Value, mlir::Value>> result;
   for (auto extent : llvm::enumerate(extents)) {
     mlir::Value lb = lowers.empty() ? one : lowers[extent.index()];
     mlir::Value ub = lowers.empty()
static mlir::Value computeVariableExtent(mlir::Location loc,
   return dimInfo.getExtent();
 }
 llvm::SmallVector<mlir::Value> getVariableExtents(mlir::Location loc,
                                                   fir::FirOpBuilder &builder,
                                                   hlfir::Entity variable) {
   llvm::SmallVector<mlir::Value> extents;
   if (fir::FortranVariableOpInterface varIface =
           variable.getIfVariableInterface()) {
-    extents = getExplicitExtents(varIface);
+    extents = getExplicitExtents(varIface, builder);
     if (!extents.empty())
       return extents;
   }
 
   if (variable.isMutableBox())
     variable = hlfir::derefPointersAndAllocatables(loc, builder, variable);
   // Use the type shape information, and/or the fir.box/fir.class shape
   // information if any extents are not static.
mlir::Value hlfir::genShape(mlir::Location loc, fir::FirOpBuilder &builder,
   // build a new fir.shape.
   return builder.create<fir::ShapeOp>(loc,
                                       getVariableExtents(loc, builder, entity));
 }
 
 llvm::SmallVector<mlir::Value>
 hlfir::getIndexExtents(mlir::Location loc, fir::FirOpBuilder &builder,
                        mlir::Value shape) {
-  llvm::SmallVector<mlir::Value> extents = getExplicitExtentsFromShape(shape);
+  llvm::SmallVector<mlir::Value> extents =
+      getExplicitExtentsFromShape(shape, builder);
   mlir::Type indexType = builder.getIndexType();
   for (auto &extent : extents)
     extent = builder.createConvert(loc, indexType, extent);
   return extents;
 }
 
 mlir::Value hlfir::genExtent(mlir::Location loc, fir::FirOpBuilder &builder,
                              hlfir::Entity entity, unsigned dim) {
   entity = followShapeInducingSource(entity);
   if (auto shape = tryRetrievingShapeOrShift(entity)) {
-    auto extents = getExplicitExtentsFromShape(shape);
+    auto extents = getExplicitExtentsFromShape(shape, builder);
     if (!extents.empty()) {
       assert(extents.size() > dim && "bad inquiry");
       return extents[dim];
     }
   }
   if (entity.isVariable()) {
     if (entity.isMutableBox())
       entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);

flang/test/HLFIR/extents-of-shape-of.f90

+! RUN: bbc -emit-fir -hlfir %s -o - | FileCheck %s
+subroutine foo(a, b)
+  real :: a(:, :), b(:, :)
+  interface
+    elemental subroutine elem_sub(x)
+      real, intent(in) :: x
+    end subroutine
+  end interface
+  call elem_sub(matmul(a, b))
+end subroutine
+! CHECK-LABEL: func.func @_QPfoo
+! CHECK:           %[[A_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {fir.bindc_name = "a"}
+! CHECK:           %[[B_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {fir.bindc_name = "b"}
+! CHECK-DAG:     %[[A_VAR:.*]]:2 = hlfir.declare %[[A_ARG]]
+! CHECK-DAG:     %[[B_VAR:.*]]:2 = hlfir.declare %[[B_ARG]]
+! CHECK-NEXT:    %[[MUL:.*]] = hlfir.matmul %[[A_VAR]]#0 %[[B_VAR]]#0
+! CHECK-NEXT:    %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr<?x?xf32>) -> !fir.shape<2>
+! CHECK-NEXT:    %[[EXT0:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 0 : index} : (!fir.shape<2>) -> index
+! CHECK-NEXT:    %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<2>) -> index
+! CHECK-NEXT:    %[[C1:.*]] = arith.constant 1 : index
+! CHECK-NEXT:    fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] {
+! CHECK-NEXT:      fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] {
+! CHECK-NEXT:        %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr<?x?xf32>, index, index) -> f32
+! CHECK-NEXT:        %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, i1)
+! CHECK-NEXT:        fir.call
+! CHECK-NEXT:        hlfir.end_associate
+! CHECK-NEXT:      }
+! CHECK-NEXT:    }
+! CHECK-NEXT:    hlfir.destroy %[[MUL]]
+! CHECK-NEXT:    return
+! CHECK-NEXT:  }