Skip to content

Commit

Permalink
[flang] lower SHAPE intrinsic (#89785)
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@jeanPerier
jeanPerier committed Apr 24, 2024
1 parent ca4dbc2 commit 3328ccf
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 0 deletions.
2 changes: 2 additions & 0 deletions flang/include/flang/Optimizer/Builder/IntrinsicCall.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>);
Expand Down
33 changes: 33 additions & 0 deletions flang/lib/Optimizer/Builder/IntrinsicCall.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>},
Expand Down Expand Up @@ -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,
Expand Down
74 changes: 74 additions & 0 deletions flang/test/Lower/Intrinsics/shape.f90
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
! 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>>)

subroutine test_kind()
interface
function return_array()
real, pointer :: return_array(:, :, :)
end function
end interface
print *, shape(return_array(), kind=8)
end subroutine
! CHECK-LABEL: func.func @_QPtest_kind() {
! CHECK: %[[VAL_0:.*]] = fir.alloca !fir.array<3xi64>
! 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) -> i64
! CHECK: %[[VAL_13:.*]] = arith.constant 0 : index
! CHECK: %[[VAL_14:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_13]] : (!fir.ref<!fir.array<3xi64>>, index) -> !fir.ref<i64>
! CHECK: fir.store %[[VAL_12]] to %[[VAL_14]] : !fir.ref<i64>
! 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) -> i64
! CHECK: %[[VAL_18:.*]] = arith.constant 1 : index
! CHECK: %[[VAL_19:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_18]] : (!fir.ref<!fir.array<3xi64>>, index) -> !fir.ref<i64>
! CHECK: fir.store %[[VAL_17]] to %[[VAL_19]] : !fir.ref<i64>
! 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) -> i64
! CHECK: %[[VAL_23:.*]] = arith.constant 2 : index
! CHECK: %[[VAL_24:.*]] = fir.coordinate_of %[[VAL_0]], %[[VAL_23]] : (!fir.ref<!fir.array<3xi64>>, index) -> !fir.ref<i64>
! CHECK: fir.store %[[VAL_22]] to %[[VAL_24]] : !fir.ref<i64>
! 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<3xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<3xi64>>, !fir.ref<!fir.array<3xi64>>)

0 comments on commit 3328ccf

Please sign in to comment.