[flang] lower SHAPE intrinsic

[jeanPerier]

Semantics usually fold SHAPE into an array constructor, but sometimes it cannot (like when the source is a function result that cannot be duplicated in expression analysis). Add lowering handling for shape.

[llvmbot]

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

Author: None (jeanPerier)

Changes

Semantics usually fold SHAPE into an array constructor, but sometimes it cannot (like when the source is a function result that cannot be duplicated in expression analysis). Add lowering handling for shape.

---

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

3 Files Affected:

• (modified) flang/include/flang/Optimizer/Builder/IntrinsicCall.h (+2)
• (modified) flang/lib/Optimizer/Builder/IntrinsicCall.cpp (+33)
• (added) flang/test/Lower/Intrinsics/shape.f90 (+38)

diff --git a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
index 6927488517e63b..604f2bd969eed5 100644
--- a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
+++ b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
@@ -335,6 +335,8 @@ struct IntrinsicLibrary {
   mlir::Value genSelectedRealKind(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSetExponent(mlir::Type resultType,
                              llvm::ArrayRef<mlir::Value> args);
+  fir::ExtendedValue genShape(mlir::Type resultType,
+                              llvm::ArrayRef<fir::ExtendedValue>);
   template <typename Shift>
   mlir::Value genShift(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
   mlir::Value genShiftA(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index 4ee7258004fa74..e28d14cd318d36 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -554,6 +554,10 @@ static constexpr IntrinsicHandler handlers[]{
        {"radix", asAddr, handleDynamicOptional}}},
      /*isElemental=*/false},
     {"set_exponent", &I::genSetExponent},
+    {"shape",
+     &I::genShape,
+     {{{"source", asBox}, {"kind", asValue}}},
+     /*isElemental=*/false},
     {"shifta", &I::genShiftA},
     {"shiftl", &I::genShift<mlir::arith::ShLIOp>},
     {"shiftr", &I::genShift<mlir::arith::ShRUIOp>},
@@ -5821,6 +5825,35 @@ mlir::Value IntrinsicLibrary::genSetExponent(mlir::Type resultType,
                                    fir::getBase(args[1])));
 }
 
+// SHAPE
+fir::ExtendedValue
+IntrinsicLibrary::genShape(mlir::Type resultType,
+                           llvm::ArrayRef<fir::ExtendedValue> args) {
+  assert(args.size() >= 1);
+  const fir::ExtendedValue &array = args[0];
+  int rank = array.rank();
+  if (rank == 0)
+    TODO(loc, "shape intrinsic lowering with assumed-rank source");
+  mlir::Type indexType = builder.getIndexType();
+  mlir::Type extentType = fir::unwrapSequenceType(resultType);
+  mlir::Type seqType = fir::SequenceType::get(
+      {static_cast<fir::SequenceType::Extent>(rank)}, extentType);
+  mlir::Value shapeArray = builder.createTemporary(loc, seqType);
+  mlir::Type shapeAddrType = builder.getRefType(extentType);
+  for (int dim = 0; dim < rank; ++dim) {
+    mlir::Value extent = fir::factory::readExtent(builder, loc, array, dim);
+    extent = builder.createConvert(loc, extentType, extent);
+    auto index = builder.createIntegerConstant(loc, indexType, dim);
+    auto shapeAddr = builder.create<fir::CoordinateOp>(loc, shapeAddrType,
+                                                       shapeArray, index);
+    builder.create<fir::StoreOp>(loc, extent, shapeAddr);
+  }
+  mlir::Value shapeArrayExtent =
+      builder.createIntegerConstant(loc, indexType, rank);
+  llvm::SmallVector<mlir::Value> extents{shapeArrayExtent};
+  return fir::ArrayBoxValue{shapeArray, extents};
+}
+
 // SHIFTL, SHIFTR
 template <typename Shift>
 mlir::Value IntrinsicLibrary::genShift(mlir::Type resultType,
diff --git a/flang/test/Lower/Intrinsics/shape.f90 b/flang/test/Lower/Intrinsics/shape.f90
new file mode 100644
index 00000000000000..97dbfda5937327
--- /dev/null
+++ b/flang/test/Lower/Intrinsics/shape.f90
@@ -0,0 +1,38 @@
+! Test SHAPE with function results
+! RUN: bbc -emit-hlfir -o - %s | FileCheck %s
+
+subroutine test()
+  interface
+    function return_array()
+      real, pointer :: return_array(:, :, :)
+    end function
+  end interface
+  print *, shape(return_array())
+end subroutine
+! CHECK-LABEL:   func.func @_QPtest() {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.array<3xi32>
+! CHECK:           %[[VAL_7:.*]] = fir.call @_QPreturn_array() {{.*}}: () -> !fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>
+! CHECK:           fir.save_result %[[VAL_7]] to %[[VAL_1:.*]] : !fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>, !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>>
+! CHECK:           %[[VAL_8:.*]]:2 = hlfir.declare %[[VAL_1]] {uniq_name = ".tmp.func_result"} : (!fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>>) -> (!fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>>, !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>>)
+! CHECK:           %[[VAL_9:.*]] = fir.load %[[VAL_8]]#1 : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>>
+! CHECK:           %[[VAL_10:.*]] = arith.constant 0 : index
+! CHECK:           %[[VAL_11:.*]]:3 = fir.box_dims %[[VAL_9]], %[[VAL_10]] : (!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>, index) -> (index, index, index)
+! CHECK:           %[[VAL_12:.*]] = fir.convert %[[VAL_11]]#1 : (index) -> i32
+! CHECK:           %[[VAL_13:.*]] = arith.constant 0 : index
+! CHECK:           %[[VAL_14:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_13]] : (!fir.ref<!fir.array<3xi32>>, index) -> !fir.ref<i32>
+! CHECK:           fir.store %[[VAL_12]] to %[[VAL_14]] : !fir.ref<i32>
+! CHECK:           %[[VAL_15:.*]] = arith.constant 1 : index
+! CHECK:           %[[VAL_16:.*]]:3 = fir.box_dims %[[VAL_9]], %[[VAL_15]] : (!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>, index) -> (index, index, index)
+! CHECK:           %[[VAL_17:.*]] = fir.convert %[[VAL_16]]#1 : (index) -> i32
+! CHECK:           %[[VAL_18:.*]] = arith.constant 1 : index
+! CHECK:           %[[VAL_19:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_18]] : (!fir.ref<!fir.array<3xi32>>, index) -> !fir.ref<i32>
+! CHECK:           fir.store %[[VAL_17]] to %[[VAL_19]] : !fir.ref<i32>
+! CHECK:           %[[VAL_20:.*]] = arith.constant 2 : index
+! CHECK:           %[[VAL_21:.*]]:3 = fir.box_dims %[[VAL_9]], %[[VAL_20]] : (!fir.box<!fir.ptr<!fir.array<?x?x?xf32>>>, index) -> (index, index, index)
+! CHECK:           %[[VAL_22:.*]] = fir.convert %[[VAL_21]]#1 : (index) -> i32
+! CHECK:           %[[VAL_23:.*]] = arith.constant 2 : index
+! CHECK:           %[[VAL_24:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_23]] : (!fir.ref<!fir.array<3xi32>>, index) -> !fir.ref<i32>
+! CHECK:           fir.store %[[VAL_22]] to %[[VAL_24]] : !fir.ref<i32>
+! CHECK:           %[[VAL_25:.*]] = arith.constant 3 : index
+! CHECK:           %[[VAL_26:.*]] = fir.shape %[[VAL_25]] : (index) -> !fir.shape<1>
+! CHECK:           %[[VAL_27:.*]]:2 = hlfir.declare %[[VAL_0]](%[[VAL_26]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<3xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<3xi32>>, !fir.ref<!fir.array<3xi32>>)

[vzakhari]

LGTM!

[tblah]

What about the optional KIND argument?

Otherwise looks good

[jeanPerier]

What about the optional KIND argument?

Thanks for pointing that out. It is taken into account via the "resultType" that was computed by semantic analysis. Still deserves a test, I added one.

[tblah]

Thanks for adding the test