Skip to content

Commit 0257f9c

Browse files
authored
[flang] lower SHAPE with assumed-rank arguments (#94812)
Allocate result statically on the stack (using max rank) and use the runtime to fill it in correctly.
1 parent d732a32 commit 0257f9c

File tree

4 files changed

+108
-2
lines changed

4 files changed

+108
-2
lines changed

flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,12 @@ mlir::Value genLboundDim(fir::FirOpBuilder &builder, mlir::Location loc,
3232
void genUbound(fir::FirOpBuilder &builder, mlir::Location loc,
3333
mlir::Value resultBox, mlir::Value array, mlir::Value kind);
3434

35+
/// Generate call to `Shape` runtime routine.
36+
/// First argument is a raw pointer to the result array storage that
37+
/// must be allocated by the caller.
38+
void genShape(fir::FirOpBuilder &builder, mlir::Location loc,
39+
mlir::Value resultAddr, mlir::Value arrayt, mlir::Value kind);
40+
3541
/// Generate call to `Size` runtime routine. This routine is a specialized
3642
/// version when the DIM argument is not specified by the user.
3743
mlir::Value genSize(fir::FirOpBuilder &builder, mlir::Location loc,

flang/lib/Optimizer/Builder/IntrinsicCall.cpp

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5992,15 +5992,45 @@ mlir::Value IntrinsicLibrary::genSetExponent(mlir::Type resultType,
59925992
fir::getBase(args[1])));
59935993
}
59945994

5995+
/// Generate runtime call to inquire about all the bounds/extents of an
5996+
/// assumed-rank array.
5997+
template <typename Func>
5998+
static fir::ExtendedValue genAssumedRankBoundInquiry(
5999+
fir::FirOpBuilder &builder, mlir::Location loc, mlir::Type resultType,
6000+
llvm::ArrayRef<fir::ExtendedValue> args, int kindPos, Func genRtCall) {
6001+
const fir::ExtendedValue &array = args[0];
6002+
// Allocate an array with the maximum rank, that is big enough to hold the
6003+
// result but still "small" (15 elements). Static size alloca make stack
6004+
// analysis/manipulation easier.
6005+
mlir::Type resultElementType = fir::unwrapSequenceType(resultType);
6006+
mlir::Type allocSeqType =
6007+
fir::SequenceType::get({Fortran::common::maxRank}, resultElementType);
6008+
mlir::Value resultStorage = builder.createTemporary(loc, allocSeqType);
6009+
mlir::Value arrayBox = builder.createBox(loc, array);
6010+
mlir::Value kind = isStaticallyAbsent(args, kindPos)
6011+
? builder.createIntegerConstant(
6012+
loc, builder.getI32Type(),
6013+
builder.getKindMap().defaultIntegerKind())
6014+
: fir::getBase(args[kindPos]);
6015+
genRtCall(builder, loc, resultStorage, arrayBox, kind);
6016+
mlir::Type baseType =
6017+
fir::ReferenceType::get(builder.getVarLenSeqTy(resultElementType));
6018+
mlir::Value resultBase = builder.createConvert(loc, baseType, resultStorage);
6019+
mlir::Value rank =
6020+
builder.create<fir::BoxRankOp>(loc, builder.getIndexType(), arrayBox);
6021+
return fir::ArrayBoxValue{resultBase, {rank}};
6022+
}
6023+
59956024
// SHAPE
59966025
fir::ExtendedValue
59976026
IntrinsicLibrary::genShape(mlir::Type resultType,
59986027
llvm::ArrayRef<fir::ExtendedValue> args) {
59996028
assert(args.size() >= 1);
60006029
const fir::ExtendedValue &array = args[0];
6030+
if (array.hasAssumedRank())
6031+
return genAssumedRankBoundInquiry(builder, loc, resultType, args,
6032+
/*kindPos=*/1, fir::runtime::genShape);
60016033
int rank = array.rank();
6002-
if (rank == 0)
6003-
TODO(loc, "shape intrinsic lowering with assumed-rank source");
60046034
mlir::Type indexType = builder.getIndexType();
60056035
mlir::Type extentType = fir::unwrapSequenceType(resultType);
60066036
mlir::Type seqType = fir::SequenceType::get(

flang/lib/Optimizer/Builder/Runtime/Inquiry.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,3 +87,17 @@ mlir::Value fir::runtime::genIsContiguous(fir::FirOpBuilder &builder,
8787
auto args = fir::runtime::createArguments(builder, loc, fTy, array);
8888
return builder.create<fir::CallOp>(loc, isContiguousFunc, args).getResult(0);
8989
}
90+
91+
void fir::runtime::genShape(fir::FirOpBuilder &builder, mlir::Location loc,
92+
mlir::Value resultAddr, mlir::Value array,
93+
mlir::Value kind) {
94+
mlir::func::FuncOp func =
95+
fir::runtime::getRuntimeFunc<mkRTKey(Shape)>(loc, builder);
96+
auto fTy = func.getFunctionType();
97+
auto sourceFile = fir::factory::locationToFilename(builder, loc);
98+
auto sourceLine =
99+
fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));
100+
auto args = fir::runtime::createArguments(
101+
builder, loc, fTy, resultAddr, array, kind, sourceFile, sourceLine);
102+
builder.create<fir::CallOp>(loc, func, args).getResult(0);
103+
}
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
! Test shape lowering for assumed-rank
2+
! RUN: bbc -emit-hlfir -o - %s -allow-assumed-rank | FileCheck %s
3+
4+
subroutine test_shape(x)
5+
real :: x(..)
6+
call takes_integer_array(shape(x))
7+
end subroutine
8+
! CHECK-LABEL: func.func @_QPtest_shape(
9+
! CHECK: %[[VAL_1:.*]] = fir.alloca !fir.array<15xi32>
10+
! CHECK: %[[VAL_4:.*]] = arith.constant 4 : i32
11+
! CHECK: %[[VAL_7:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.llvm_ptr<i8>
12+
! CHECK: %[[VAL_8:.*]] = fir.convert %[[VAL_3:.*]] : (!fir.box<!fir.array<*:f32>>) -> !fir.box<none>
13+
! CHECK: %[[VAL_10:.*]] = fir.call @_FortranAShape(%[[VAL_7]], %[[VAL_8]], %[[VAL_4]], %{{.*}}, %{{.*}})
14+
! CHECK: %[[VAL_11:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.ref<!fir.array<?xi32>>
15+
! CHECK: %[[VAL_12:.*]] = fir.box_rank %[[VAL_3]] : (!fir.box<!fir.array<*:f32>>) -> index
16+
! CHECK: %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1>
17+
! CHECK: %[[VAL_14:.*]]:2 = hlfir.declare %[[VAL_11]](%[[VAL_13]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>)
18+
! CHECK: %[[VAL_15:.*]] = arith.constant false
19+
! CHECK: %[[VAL_16:.*]] = hlfir.as_expr %[[VAL_14]]#0 move %[[VAL_15]] : (!fir.box<!fir.array<?xi32>>, i1) -> !hlfir.expr<?xi32>
20+
! CHECK: %[[VAL_17:.*]]:3 = hlfir.associate %[[VAL_16]](%[[VAL_13]]) {adapt.valuebyref} : (!hlfir.expr<?xi32>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>, i1)
21+
! CHECK: fir.call @_QPtakes_integer_array(%[[VAL_17]]#1) fastmath<contract> : (!fir.ref<!fir.array<?xi32>>) -> ()
22+
! CHECK: hlfir.end_associate %[[VAL_17]]#1, %[[VAL_17]]#2 : !fir.ref<!fir.array<?xi32>>, i1
23+
! CHECK: hlfir.destroy %[[VAL_16]] : !hlfir.expr<?xi32>
24+
! CHECK: return
25+
! CHECK: }
26+
27+
subroutine test_shape_kind(x)
28+
real :: x(..)
29+
call takes_integer8_array(shape(x, kind=8))
30+
end subroutine
31+
! CHECK-LABEL: func.func @_QPtest_shape_kind(
32+
! CHECK: %[[VAL_1:.*]] = fir.alloca !fir.array<15xi64>
33+
! CHECK: %[[VAL_4:.*]] = arith.constant 8 : i32
34+
! CHECK: %[[VAL_7:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi64>>) -> !fir.llvm_ptr<i8>
35+
! CHECK: %[[VAL_8:.*]] = fir.convert %[[VAL_3:.*]] : (!fir.box<!fir.array<*:f32>>) -> !fir.box<none>
36+
! CHECK: %[[VAL_10:.*]] = fir.call @_FortranAShape(%[[VAL_7]], %[[VAL_8]], %[[VAL_4]], %{{.*}}, %{{.*}})
37+
! CHECK: %[[VAL_11:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi64>>) -> !fir.ref<!fir.array<?xi64>>
38+
! CHECK: %[[VAL_12:.*]] = fir.box_rank %[[VAL_3]] : (!fir.box<!fir.array<*:f32>>) -> index
39+
! CHECK: %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1>
40+
! CHECK: %[[VAL_14:.*]]:2 = hlfir.declare %[[VAL_11]](%[[VAL_13]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi64>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi64>>, !fir.ref<!fir.array<?xi64>>)
41+
42+
subroutine test_shape_2(x)
43+
real, pointer :: x(..)
44+
call takes_integer_array(shape(x))
45+
end subroutine
46+
! CHECK-LABEL: func.func @_QPtest_shape_2(
47+
! CHECK: %[[VAL_1:.*]] = fir.alloca !fir.array<15xi32>
48+
! CHECK: %[[VAL_4:.*]] = fir.load %[[VAL_3:.*]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<*:f32>>>>
49+
! CHECK: %[[VAL_5:.*]] = arith.constant 4 : i32
50+
! CHECK: %[[VAL_8:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.llvm_ptr<i8>
51+
! CHECK: %[[VAL_9:.*]] = fir.convert %[[VAL_4]] : (!fir.box<!fir.ptr<!fir.array<*:f32>>>) -> !fir.box<none>
52+
! CHECK: %[[VAL_11:.*]] = fir.call @_FortranAShape(%[[VAL_8]], %[[VAL_9]], %[[VAL_5]], %{{.*}}, %{{.*}})
53+
! CHECK: %[[VAL_12:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<15xi32>>) -> !fir.ref<!fir.array<?xi32>>
54+
! CHECK: %[[VAL_13:.*]] = fir.box_rank %[[VAL_4]] : (!fir.box<!fir.ptr<!fir.array<*:f32>>>) -> index
55+
! CHECK: %[[VAL_14:.*]] = fir.shape %[[VAL_13]] : (index) -> !fir.shape<1>
56+
! CHECK: %[[VAL_15:.*]]:2 = hlfir.declare %[[VAL_12]](%[[VAL_14]]) {uniq_name = ".tmp.intrinsic_result"} : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ref<!fir.array<?xi32>>)

0 commit comments

Comments
 (0)