Skip to content

Commit eea7c93

Browse files
clementvaljeanPeriermleair
committed
[flang] Lower array related intrinsics
This patch adds lowering for somw array related intrinsics: - `reshape` - `spread` This patch is part of the upstreaming effort from fir-dev branch. Reviewed By: schweitz Differential Revision: https://reviews.llvm.org/D121841 Co-authored-by: Jean Perier <jperier@nvidia.com> Co-authored-by: mleair <leairmark@gmail.com>
1 parent cae735f commit eea7c93

File tree

3 files changed

+186
-0
lines changed

3 files changed

+186
-0
lines changed

‎flang/lib/Lower/IntrinsicCall.cpp

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,11 +486,13 @@ struct IntrinsicLibrary {
486486
void genRandomInit(llvm::ArrayRef<fir::ExtendedValue>);
487487
void genRandomNumber(llvm::ArrayRef<fir::ExtendedValue>);
488488
void genRandomSeed(llvm::ArrayRef<fir::ExtendedValue>);
489+
fir::ExtendedValue genReshape(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
489490
fir::ExtendedValue genScan(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
490491
mlir::Value genSetExponent(mlir::Type resultType,
491492
llvm::ArrayRef<mlir::Value> args);
492493
fir::ExtendedValue genSize(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
493494
fir::ExtendedValue genSum(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
495+
fir::ExtendedValue genSpread(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
494496
void genSystemClock(llvm::ArrayRef<fir::ExtendedValue>);
495497
fir::ExtendedValue genTransfer(mlir::Type,
496498
llvm::ArrayRef<fir::ExtendedValue>);
@@ -729,6 +731,13 @@ static constexpr IntrinsicHandler handlers[]{
729731
&I::genRandomSeed,
730732
{{{"size", asBox}, {"put", asBox}, {"get", asBox}}},
731733
/*isElemental=*/false},
734+
{"reshape",
735+
&I::genReshape,
736+
{{{"source", asBox},
737+
{"shape", asBox},
738+
{"pad", asBox, handleDynamicOptional},
739+
{"order", asBox, handleDynamicOptional}}},
740+
/*isElemental=*/false},
732741
{"scan",
733742
&I::genScan,
734743
{{{"string", asAddr},
@@ -743,6 +752,10 @@ static constexpr IntrinsicHandler handlers[]{
743752
{"dim", asAddr, handleDynamicOptional},
744753
{"kind", asValue}}},
745754
/*isElemental=*/false},
755+
{"spread",
756+
&I::genSpread,
757+
{{{"source", asAddr}, {"dim", asValue}, {"ncopies", asValue}}},
758+
/*isElemental=*/false},
746759
{"sum",
747760
&I::genSum,
748761
{{{"array", asBox},
@@ -2501,6 +2514,52 @@ void IntrinsicLibrary::genRandomSeed(llvm::ArrayRef<fir::ExtendedValue> args) {
25012514
Fortran::lower::genRandomSeed(builder, loc, -1, mlir::Value{});
25022515
}
25032516

2517+
// RESHAPE
2518+
fir::ExtendedValue
2519+
IntrinsicLibrary::genReshape(mlir::Type resultType,
2520+
llvm::ArrayRef<fir::ExtendedValue> args) {
2521+
assert(args.size() == 4);
2522+
2523+
// Handle source argument
2524+
mlir::Value source = builder.createBox(loc, args[0]);
2525+
2526+
// Handle shape argument
2527+
mlir::Value shape = builder.createBox(loc, args[1]);
2528+
assert(fir::BoxValue(shape).rank() == 1);
2529+
mlir::Type shapeTy = shape.getType();
2530+
mlir::Type shapeArrTy = fir::dyn_cast_ptrOrBoxEleTy(shapeTy);
2531+
auto resultRank = shapeArrTy.cast<fir::SequenceType>().getShape();
2532+
2533+
assert(resultRank[0] != fir::SequenceType::getUnknownExtent() &&
2534+
"shape arg must have constant size");
2535+
2536+
// Handle optional pad argument
2537+
mlir::Value pad = isAbsent(args[2])
2538+
? builder.create<fir::AbsentOp>(
2539+
loc, fir::BoxType::get(builder.getI1Type()))
2540+
: builder.createBox(loc, args[2]);
2541+
2542+
// Handle optional order argument
2543+
mlir::Value order = isAbsent(args[3])
2544+
? builder.create<fir::AbsentOp>(
2545+
loc, fir::BoxType::get(builder.getI1Type()))
2546+
: builder.createBox(loc, args[3]);
2547+
2548+
// Create mutable fir.box to be passed to the runtime for the result.
2549+
mlir::Type type = builder.getVarLenSeqTy(resultType, resultRank[0]);
2550+
fir::MutableBoxValue resultMutableBox =
2551+
fir::factory::createTempMutableBox(builder, loc, type);
2552+
2553+
mlir::Value resultIrBox =
2554+
fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
2555+
2556+
fir::runtime::genReshape(builder, loc, resultIrBox, source, shape, pad,
2557+
order);
2558+
2559+
return readAndAddCleanUp(resultMutableBox, resultType,
2560+
"unexpected result for RESHAPE");
2561+
}
2562+
25042563
// SCAN
25052564
fir::ExtendedValue
25062565
IntrinsicLibrary::genScan(mlir::Type resultType,
@@ -2589,6 +2648,38 @@ mlir::Value IntrinsicLibrary::genSetExponent(mlir::Type resultType,
25892648
fir::getBase(args[1])));
25902649
}
25912650

2651+
// SPREAD
2652+
fir::ExtendedValue
2653+
IntrinsicLibrary::genSpread(mlir::Type resultType,
2654+
llvm::ArrayRef<fir::ExtendedValue> args) {
2655+
2656+
assert(args.size() == 3);
2657+
2658+
// Handle source argument
2659+
mlir::Value source = builder.createBox(loc, args[0]);
2660+
fir::BoxValue sourceTmp = source;
2661+
unsigned sourceRank = sourceTmp.rank();
2662+
2663+
// Handle Dim argument
2664+
mlir::Value dim = fir::getBase(args[1]);
2665+
2666+
// Handle ncopies argument
2667+
mlir::Value ncopies = fir::getBase(args[2]);
2668+
2669+
// Generate result descriptor
2670+
mlir::Type resultArrayType =
2671+
builder.getVarLenSeqTy(resultType, sourceRank + 1);
2672+
fir::MutableBoxValue resultMutableBox =
2673+
fir::factory::createTempMutableBox(builder, loc, resultArrayType);
2674+
mlir::Value resultIrBox =
2675+
fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
2676+
2677+
fir::runtime::genSpread(builder, loc, resultIrBox, source, dim, ncopies);
2678+
2679+
return readAndAddCleanUp(resultMutableBox, resultType,
2680+
"unexpected result for SPREAD");
2681+
}
2682+
25922683
// SUM
25932684
fir::ExtendedValue
25942685
IntrinsicLibrary::genSum(mlir::Type resultType,

‎flang/test/Lower/Intrinsics/reshape.f90

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
! RUN: bbc -emit-fir %s -o - | FileCheck %s
2+
3+
! CHECK-LABEL: func @_QPreshape_test(
4+
! CHECK-SAME: %[[arg0:.*]]: !fir.box<!fir.array<?x?xi32>>{{.*}}, %[[arg1:[^:]+]]: !fir.box<!fir.array<?x?x?xi32>>{{.*}}, %[[arg2:[^:]+]]: !fir.box<!fir.array<?x?x?xi32>>{{.*}}, %[[arg3:.*]]: !fir.ref<!fir.array<2xi32>>{{.*}}, %[[arg4:.*]]: !fir.ref<!fir.array<2xi32>>{{.*}}) {
5+
subroutine reshape_test(x, source, pd, sh, ord)
6+
integer :: x(:,:)
7+
integer :: source(:,:,:)
8+
integer :: pd(:,:,:)
9+
integer :: sh(2)
10+
integer :: ord(2)
11+
! CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
12+
! CHECK-DAG: %[[a0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xi32>>>
13+
! CHECK-DAG: %[[a1:.*]] = fir.shape %[[c2]] : (index) -> !fir.shape<1>
14+
! CHECK-DAG: %[[a2:.*]] = fir.embox %[[arg3]](%{{.*}}) : (!fir.ref<!fir.array<2xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<2xi32>>
15+
! CHECK-DAG: %[[a3:.*]] = fir.embox %[[arg4]](%{{.*}}) : (!fir.ref<!fir.array<2xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<2xi32>>
16+
! CHECK-DAG: %[[a8:.*]] = fir.convert %[[a0]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>) -> !fir.ref<!fir.box<none>>
17+
! CHECK-DAG: %[[a9:.*]] = fir.convert %[[arg1]] : (!fir.box<!fir.array<?x?x?xi32>>) -> !fir.box<none>
18+
! CHECK-DAG: %[[a10:.*]] = fir.convert %[[a2]] : (!fir.box<!fir.array<2xi32>>) -> !fir.box<none>
19+
! CHECK-DAG: %[[a11:.*]] = fir.convert %[[arg2]] : (!fir.box<!fir.array<?x?x?xi32>>) -> !fir.box<none>
20+
! CHECK-DAG: %[[a12:.*]] = fir.convert %[[a3]] : (!fir.box<!fir.array<2xi32>>) -> !fir.box<none>
21+
x = reshape(source, sh, pd, ord)
22+
! CHECK: %{{.*}} = fir.call @_FortranAReshape(%[[a8]], %[[a9]], %[[a10]], %[[a11]], %[[a12]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> none
23+
! CHECK-DAG: %[[a15:.*]] = fir.load %[[a0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>
24+
! CHECK-DAG: %[[a18:.*]] = fir.box_addr %[[a15]] : (!fir.box<!fir.heap<!fir.array<?x?xi32>>>) -> !fir.heap<!fir.array<?x?xi32>>
25+
! CHECK-DAG: fir.freemem %[[a18]]
26+
end subroutine
27+
28+
! CHECK-LABEL: func @_QPtest_reshape_optional(
29+
! CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
30+
! CHECK-SAME: %[[VAL_1:.*]]: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
31+
subroutine test_reshape_optional(pad, order, source, shape)
32+
real, pointer :: pad(:, :)
33+
integer, pointer :: order(:)
34+
real :: source(:, :, :)
35+
integer :: shape(4)
36+
print *, reshape(source=source, shape=shape, pad=pad, order=order)
37+
! CHECK: %[[VAL_13:.*]] = fir.load %[[VAL_0]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
38+
! CHECK: %[[VAL_14:.*]] = fir.box_addr %[[VAL_13]] : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>) -> !fir.ptr<!fir.array<?x?xf32>>
39+
! CHECK: %[[VAL_15:.*]] = fir.convert %[[VAL_14]] : (!fir.ptr<!fir.array<?x?xf32>>) -> i64
40+
! CHECK: %[[VAL_16:.*]] = arith.constant 0 : i64
41+
! CHECK: %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_15]], %[[VAL_16]] : i64
42+
! CHECK: %[[VAL_18:.*]] = fir.load %[[VAL_0]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?x?xf32>>>>
43+
! CHECK: %[[VAL_19:.*]] = fir.absent !fir.box<!fir.ptr<!fir.array<?x?xf32>>>
44+
! CHECK: %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_18]], %[[VAL_19]] : !fir.box<!fir.ptr<!fir.array<?x?xf32>>>
45+
! CHECK: %[[VAL_21:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
46+
! CHECK: %[[VAL_22:.*]] = fir.box_addr %[[VAL_21]] : (!fir.box<!fir.ptr<!fir.array<?xi32>>>) -> !fir.ptr<!fir.array<?xi32>>
47+
! CHECK: %[[VAL_23:.*]] = fir.convert %[[VAL_22]] : (!fir.ptr<!fir.array<?xi32>>) -> i64
48+
! CHECK: %[[VAL_24:.*]] = arith.constant 0 : i64
49+
! CHECK: %[[VAL_25:.*]] = arith.cmpi ne, %[[VAL_23]], %[[VAL_24]] : i64
50+
! CHECK: %[[VAL_26:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>
51+
! CHECK: %[[VAL_27:.*]] = fir.absent !fir.box<!fir.ptr<!fir.array<?xi32>>>
52+
! CHECK: %[[VAL_28:.*]] = arith.select %[[VAL_25]], %[[VAL_26]], %[[VAL_27]] : !fir.box<!fir.ptr<!fir.array<?xi32>>>
53+
! CHECK: %[[VAL_38:.*]] = fir.convert %[[VAL_20]] : (!fir.box<!fir.ptr<!fir.array<?x?xf32>>>) -> !fir.box<none>
54+
! CHECK: %[[VAL_39:.*]] = fir.convert %[[VAL_28]] : (!fir.box<!fir.ptr<!fir.array<?xi32>>>) -> !fir.box<none>
55+
! CHECK: %[[VAL_41:.*]] = fir.call @_FortranAReshape({{.*}}, {{.*}}, %{{.*}}, %[[VAL_38]], %[[VAL_39]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> none
56+
end subroutine

‎flang/test/Lower/Intrinsics/spread.f90

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
! RUN: bbc -emit-fir %s -o - | FileCheck %s
2+
3+
! CHECK-LABEL: func @_QPspread_test(
4+
! CHECK-SAME: %[[arg0:[^:]+]]: !fir.ref<i32>{{.*}}, %[[arg1:[^:]+]]: !fir.ref<i32>{{.*}}, %[[arg2:[^:]+]]: !fir.ref<i32>{{.*}}, %[[arg3:.*]]: !fir.box<!fir.array<?xi32>>{{.*}}) {
5+
subroutine spread_test(s,d,n,r)
6+
integer :: s,d,n
7+
integer :: r(:)
8+
! CHECK-DAG: %[[a0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>>
9+
! CHECK-DAG: %[[a1:.*]] = fir.load %[[arg1]] : !fir.ref<i32>
10+
! CHECK-DAG: %[[a2:.*]] = fir.load %[[arg2]] : !fir.ref<i32>
11+
! CHECK-DAG: %[[a3:.*]] = fir.embox %[[arg0]] : (!fir.ref<i32>) -> !fir.box<i32>
12+
! CHECK-DAG: %[[a8:.*]] = fir.convert %[[a0]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<none>>
13+
! CHECK-DAG: %[[a9:.*]] = fir.convert %[[a3]] : (!fir.box<i32>) -> !fir.box<none>
14+
! CHECK-DAG: %[[a10:.*]] = fir.convert %[[a2]] : (i32) -> i64
15+
r = spread(s,d,n)
16+
! CHECK: %{{.*}} = fir.call @_FortranASpread(%[[a8]], %[[a9]], %[[a1]], %[[a10]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i64, !fir.ref<i8>, i32) -> none
17+
! CHECK-DAG: %[[a13:.*]] = fir.load %[[a0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
18+
! CHECK-DAG: %[[a15:.*]] = fir.box_addr %[[a13]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
19+
! CHECK: fir.freemem %[[a15]]
20+
end subroutine
21+
22+
! CHECK-LABEL: func @_QPspread_test2(
23+
! CHECK-SAME: %[[arg0:.*]]: !fir.box<!fir.array<?xi32>>{{.*}}, %[[arg1:[^:]+]]: !fir.ref<i32>{{.*}}, %[[arg2:[^:]+]]: !fir.ref<i32>{{.*}}, %[[arg3:.*]]: !fir.box<!fir.array<?x?xi32>>{{.*}}) {
24+
subroutine spread_test2(s,d,n,r)
25+
integer :: s(:),d,n
26+
integer :: r(:,:)
27+
! CHECK-DAG: %[[a0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xi32>>>
28+
! CHECK-DAG: %[[a1:.*]] = fir.load %[[arg1]] : !fir.ref<i32>
29+
! CHECK-DAG: %[[a2:.*]] = fir.load %[[arg2]] : !fir.ref<i32>
30+
! CHECK-DAG: %[[a7:.*]] = fir.convert %[[a0]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>) -> !fir.ref<!fir.box<none>>
31+
! CHECK-DAG: %[[a8:.*]] = fir.convert %[[arg0]] : (!fir.box<!fir.array<?xi32>>) -> !fir.box<none>
32+
! CHECK-DAG: %[[a9:.*]] = fir.convert %[[a2]] : (i32) -> i64
33+
r = spread(s,d,n)
34+
! CHECK: %{{.*}} = fir.call @_FortranASpread(%[[a7]], %[[a8]], %[[a1]], %[[a9]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i64, !fir.ref<i8>, i32) -> none
35+
! CHECK-DAG: %[[a12:.*]] = fir.load %[[a0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>
36+
! CHECK-DAG: %[[a15:.*]] = fir.box_addr %[[a12]] : (!fir.box<!fir.heap<!fir.array<?x?xi32>>>) -> !fir.heap<!fir.array<?x?xi32>>
37+
! CHECK: fir.freemem %[[a15:.*]]
38+
end subroutine
39+

0 commit comments

Comments
 (0)