[flang][hlfir] Lower evaluate::DescriptorInquiry

[jeanPerier]

Summary

Lower extents and lower bounds inquiries with a compile time
constant DIM that cannot be folded by the front-end.

[clementval]

LGTM

flang/test/Lower/HLFIR/descriptor-inquiries.f90

3

flang/include/flang/Optimizer/Builder/HLFIRTools.h

 /// (fir.shift is not allowed here).
 llvm::SmallVector<std::pair<mlir::Value, mlir::Value>>
 genBounds(mlir::Location loc, fir::FirOpBuilder &builder, mlir::Value shape);
 
 /// Compute fir.shape<> (no lower bounds) for an entity.
 mlir::Value genShape(mlir::Location loc, fir::FirOpBuilder &builder,
                      Entity entity);
 
+/// Compute the extent of \p entity in dimension \p dim. Crashes
+/// if dim is bigger than the entity's rank.
+mlir::Value genExtent(mlir::Location loc, fir::FirOpBuilder &builder,
+                      hlfir::Entity entity, unsigned dim);
+
+/// Compute the lower bound of \p entity in dimension \p dim. Crashes
+/// if dim is bigger than the entity's rank.
+mlir::Value genLBound(mlir::Location loc, fir::FirOpBuilder &builder,
+                      hlfir::Entity entity, unsigned dim);
+
 /// Generate a vector of extents with index type from a fir.shape
 /// of fir.shape_shift value.
 llvm::SmallVector<mlir::Value> getIndexExtents(mlir::Location loc,
                                                fir::FirOpBuilder &builder,
                                                mlir::Value shape);
 
 /// Read length parameters into result if this entity has any.
 void genLengthParameters(mlir::Location loc, fir::FirOpBuilder &builder,

flang/lib/Lower/ConvertExprToHLFIR.cpp

gen(const Fortran::evaluate::DescriptorInquiry &desc) {
     auto castResult = [&](mlir::Value v) {
       return hlfir::EntityWithAttributes{
           builder.createConvert(loc, resultType, v)};
     };
     switch (desc.field()) {
     case Fortran::evaluate::DescriptorInquiry::Field::Len:
       return castResult(hlfir::genCharLength(loc, builder, entity));
     case Fortran::evaluate::DescriptorInquiry::Field::LowerBound:
-      TODO(loc, "lower bound inquiry in HLFIR");
+      return castResult(
+          hlfir::genLBound(loc, builder, entity, desc.dimension()));
     case Fortran::evaluate::DescriptorInquiry::Field::Extent:
-      TODO(loc, "extent inquiry in HLFIR");
+      return castResult(
+          hlfir::genExtent(loc, builder, entity, desc.dimension()));
     case Fortran::evaluate::DescriptorInquiry::Field::Rank:
       TODO(loc, "rank inquiry on assumed rank");
     case Fortran::evaluate::DescriptorInquiry::Field::Stride:
       // So far the front end does not generate this inquiry.
       TODO(loc, "stride inquiry");
     }
     llvm_unreachable("unknown descriptor inquiry");
   }

flang/lib/Optimizer/Builder/HLFIRTools.cpp

for (auto extent : llvm::enumerate(extents)) {
     mlir::Value ub = lowers.empty()
                          ? extent.value()
                          : genUBound(loc, builder, lb, extent.value(), one);
     result.push_back({lb, ub});
   }
   return result;
 }
 
-static hlfir::Entity followEntitySource(hlfir::Entity entity) {
+static hlfir::Entity followShapeInducingSource(hlfir::Entity entity) {
   while (true) {
     if (auto reassoc = entity.getDefiningOp<hlfir::NoReassocOp>()) {
       entity = hlfir::Entity{reassoc.getVal()};
       continue;
     }
     if (auto asExpr = entity.getDefiningOp<hlfir::AsExprOp>()) {
       entity = hlfir::Entity{asExpr.getVar()};
       continue;
     }
     break;
   }
   return entity;
 }
 
+static mlir::Value computeVariableExtent(mlir::Location loc,
+                                         fir::FirOpBuilder &builder,
+                                         hlfir::Entity variable,
+                                         fir::SequenceType seqTy,
+                                         unsigned dim) {
+  mlir::Type idxTy = builder.getIndexType();
+  if (seqTy.getShape().size() > dim) {
+    fir::SequenceType::Extent typeExtent = seqTy.getShape()[dim];
+    if (typeExtent != fir::SequenceType::getUnknownExtent())
+      return builder.createIntegerConstant(loc, idxTy, typeExtent);
+  }
+  assert(variable.getType().isa<fir::BaseBoxType>() &&
+         "array variable with dynamic extent must be boxed");
+  mlir::Value dimVal = builder.createIntegerConstant(loc, idxTy, dim);
+  auto dimInfo = builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy,
+                                                variable, dimVal);
+  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);
     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.
   fir::SequenceType seqTy =
       hlfir::getFortranElementOrSequenceType(variable.getType())
           .cast<fir::SequenceType>();
-  mlir::Type idxTy = builder.getIndexType();
-  for (auto typeExtent : seqTy.getShape())
-    if (typeExtent != fir::SequenceType::getUnknownExtent()) {
-      extents.push_back(builder.createIntegerConstant(loc, idxTy, typeExtent));
-    } else {
-      assert(variable.getType().isa<fir::BaseBoxType>() &&
-             "array variable with dynamic extent must be boxed");
-      mlir::Value dim =
-          builder.createIntegerConstant(loc, idxTy, extents.size());
-      auto dimInfo = builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy,
-                                                    variable, dim);
-      extents.push_back(dimInfo.getExtent());
-    }
+  unsigned rank = seqTy.getShape().size();
+  for (unsigned dim = 0; dim < rank; ++dim)
+    extents.push_back(
+        computeVariableExtent(loc, builder, variable, seqTy, dim));
   return extents;
 }
 
-mlir::Value hlfir::genShape(mlir::Location loc, fir::FirOpBuilder &builder,
-                            hlfir::Entity entity) {
-  assert(entity.isArray() && "entity must be an array");
-  entity = followEntitySource(entity);
-
+static mlir::Value tryRetrievingShapeOrShift(hlfir::Entity entity) {
   if (entity.getType().isa<hlfir::ExprType>()) {
     if (auto elemental = entity.getDefiningOp<hlfir::ElementalOp>())
       return elemental.getShape();
-    TODO(loc, "get shape from HLFIR expr without producer holding the shape");
+    return mlir::Value{};
   }
-  // Entity is an array variable.
-  if (auto varIface = entity.getIfVariableInterface()) {
-    if (auto shape = varIface.getShape()) {
-      if (shape.getType().isa<fir::ShapeType>())
-        return shape;
-      if (shape.getType().isa<fir::ShapeShiftType>())
-        if (auto s = shape.getDefiningOp<fir::ShapeShiftOp>())
-          return builder.create<fir::ShapeOp>(loc, s.getExtents());
-    }
-  }
+  if (auto varIface = entity.getIfVariableInterface())
+    return varIface.getShape();
+  return {};
+}
+
+mlir::Value hlfir::genShape(mlir::Location loc, fir::FirOpBuilder &builder,
+                            hlfir::Entity entity) {
+  assert(entity.isArray() && "entity must be an array");
+  entity = followShapeInducingSource(entity);
+  assert(entity && "what?");
+  if (auto shape = tryRetrievingShapeOrShift(entity)) {
+  if (entity.getType().isa<hlfir::ExprType>())
+    TODO(loc, "get shape from HLFIR expr without producer holding the shape");
   // There is no shape lying around for this entity. Retrieve the extents and
   // 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);
   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);
+    if (!extents.empty()) {
+      assert(extents.size() > dim && "bad inquiry");
+      return extents[dim];
+    }
+  }
+  if (entity.isVariable()) {
+    if (entity.isMutableBox())
+      entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
+    // Use the type shape information, and/or the fir.box/fir.class shape
+    // information if any extents are not static.
+    fir::SequenceType seqTy =
+        hlfir::getFortranElementOrSequenceType(entity.getType())
+            .cast<fir::SequenceType>();
+    return computeVariableExtent(loc, builder, entity, seqTy, dim);
+  }
+  TODO(loc, "get extent from HLFIR expr without producer holding the shape");
+}
+
+mlir::Value hlfir::genLBound(mlir::Location loc, fir::FirOpBuilder &builder,
+                             hlfir::Entity entity, unsigned dim) {
+  if (!entity.hasNonDefaultLowerBounds())
+    return builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+  if (auto shape = tryRetrievingShapeOrShift(entity)) {
+    auto lbounds = getExplicitLboundsFromShape(shape);
+    if (!lbounds.empty()) {
+      assert(lbounds.size() > dim && "bad inquiry");
+      return lbounds[dim];
+    }
+  }
+  if (entity.isMutableBox())
+    entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
+  assert(entity.getType().isa<fir::BaseBoxType>() && "must be a box");
+  mlir::Type idxTy = builder.getIndexType();
+  mlir::Value dimVal = builder.createIntegerConstant(loc, idxTy, dim);
+  auto dimInfo =
+      builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy, entity, dimVal);
+  return dimInfo.getLowerBound();
+}
+
 void hlfir::genLengthParameters(mlir::Location loc, fir::FirOpBuilder &builder,
                                 Entity entity,
                                 llvm::SmallVectorImpl<mlir::Value> &result) {
   if (!entity.hasLengthParameters())
     return;
   if (entity.getType().isa<hlfir::ExprType>()) {
     mlir::Value expr = entity;
     if (auto reassoc = expr.getDefiningOp<hlfir::NoReassocOp>())

flang/test/Lower/HLFIR/descriptor-inquiries.f90

+! Test lowering of extent and lower bound inquires that
+! come in lowering as evaluate::DescriptorInquiry.
+
+! RUN: bbc -emit-fir -hlfir -o - %s | FileCheck %s
+
+subroutine test_assumed_shape(x, r)
+  integer(8) :: r
+  real :: x(:,:)
+  r = size(x, dim=2, kind=8)
+end subroutine
+! CHECK-LABEL: func.func @_QPtest_assumed_shape(
+! CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare {{.*}}Ex
+! CHECK:  %[[VAL_4:.*]] = arith.constant 1 : index
+! CHECK:  %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_3]]#0, %[[VAL_4]] : (!fir.box<!fir.array<?x?xf32>>, index) -> (index, index, index)
+! CHECK:  %[[VAL_6:.*]] = fir.convert %[[VAL_5]]#1 : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_6]] to %{{.*}}
+
+subroutine test_explicit_shape(x, n, m, r)
+  integer(8) :: n, m, r
+  real :: x(n,m)
+  r = size(x, dim=2, kind=8)
+end subroutine
+! CHECK-LABEL: func.func @_QPtest_explicit_shape(
+! CHECK:  %[[VAL_17:.*]] = fir.shape %{{.*}}, %[[VAL_16:.*]] : (index, index) -> !fir.shape<2>
+! CHECK:  %[[VAL_18:.*]]:2 = hlfir.declare %{{.*}}(%[[VAL_17]])  {{.*}}Ex
+! CHECK:  %[[VAL_19:.*]] = fir.convert %[[VAL_16]] : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_19]] to %{{.*}}
+
+subroutine test_pointer(x, r)
+  integer(8) :: r
+  real :: x(:,:)
+  r = size(x, dim=2, kind=8)
+end subroutine
+! CHECK-LABEL: func.func @_QPtest_pointer(
+! CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare {{.*}}Ex
+! CHECK:  %[[VAL_4:.*]] = arith.constant 1 : index
+! CHECK:  %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_3]]#0, %[[VAL_4]] : (!fir.box<!fir.array<?x?xf32>>, index) -> (index, index, index)
+! CHECK:  %[[VAL_6:.*]] = fir.convert %[[VAL_5]]#1 : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_6]] to %{{.*}}
+
+subroutine test_lbound_assumed_shape(x, l1, l2, r)
+  integer(8) :: l1, l2, r
+  real :: x(l1:,l2:)
+  r = lbound(x, dim=2, kind=8)
+end subroutine
+! CHECK:  %[[VAL_11:.*]] = fir.shift %[[VAL_8:.*]], %[[VAL_10:.*]] : (index, index) -> !fir.shift<2>
+! CHECK:  %[[VAL_12:.*]]:2 = hlfir.declare %{{.*}}(%[[VAL_11]])  {{.*}}Ex
+! CHECK:  %[[VAL_13:.*]] = arith.constant 1 : i64
+! CHECK:  %[[VAL_14:.*]] = arith.constant 0 : index
+! CHECK:  %[[VAL_15:.*]] = arith.constant 1 : index
+! CHECK:  %[[VAL_16:.*]]:3 = fir.box_dims %[[VAL_12]]#1, %[[VAL_15]] : (!fir.box<!fir.array<?x?xf32>>, index) -> (index, index, index)
+! CHECK:  %[[VAL_17:.*]] = arith.cmpi eq, %[[VAL_16]]#1, %[[VAL_14]] : index
+! CHECK:  %[[VAL_18:.*]] = fir.convert %[[VAL_13]] : (i64) -> index
+! CHECK:  %[[VAL_19:.*]] = arith.select %[[VAL_17]], %[[VAL_18]], %[[VAL_10]] : index
+! CHECK:  %[[VAL_20:.*]] = fir.convert %[[VAL_19]] : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_20]] to %{{.*}}
+
+subroutine test_lbound_explicit_shape(x, n, m, l1, l2, r)
+  integer(8) :: n, m, l1, l2, r
+  real :: x(l1:n,l2:m)
+  r = lbound(x, dim=2, kind=8)
+end subroutine
+! CHECK-LABEL: func.func @_QPtest_lbound_explicit_shape(
+! CHECK:  %[[VAL_31:.*]] = fir.shape_shift %{{.*}}, %{{.*}}, %[[VAL_22:.*]], %[[VAL_30:.*]] : (index, index, index, index) -> !fir.shapeshift<2>
+! CHECK:  %[[VAL_32:.*]]:2 = hlfir.declare %{{.*}}(%[[VAL_31]])  {{.*}}Ex
+! CHECK:  %[[VAL_33:.*]] = arith.constant 1 : i64
+! CHECK:  %[[VAL_34:.*]] = arith.constant 0 : index
+! CHECK:  %[[VAL_35:.*]] = arith.cmpi eq, %[[VAL_30]], %[[VAL_34]] : index
+! CHECK:  %[[VAL_36:.*]] = fir.convert %[[VAL_33]] : (i64) -> index
+! CHECK:  %[[VAL_37:.*]] = arith.select %[[VAL_35]], %[[VAL_36]], %[[VAL_22]] : index
+! CHECK:  %[[VAL_38:.*]] = fir.convert %[[VAL_37]] : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_38]] to %{{.*}}
+
+subroutine test_lbound_pointer(x, r)
+  integer(8) :: r
+  real, pointer :: x(:,:)
+  r = lbound(x, dim=2, kind=8)
+end subroutine
+! CHECK-LABEL: func.func @_QPtest_lbound_pointer(
+! CHECK:  %[[VAL_3:.*]]:2 = hlfir.declare {{.*}}Ex
+! CHECK:  %[[VAL_4:.*]] = fir.load %[[VAL_3]]#0 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
+! CHECK:  %[[VAL_5:.*]] = arith.constant 1 : index
+! CHECK:  %[[VAL_6:.*]]:3 = fir.box_dims %[[VAL_4]], %[[VAL_5]] : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>, index) -> (index, index, index)
+! CHECK:  %[[VAL_7:.*]] = fir.convert %[[VAL_6]]#0 : (index) -> i64
+! CHECK:  hlfir.assign %[[VAL_7]] to %{{.*}}