Skip to content

Commit

Permalink
[flang] Lower F08 FINDLOC intrinsic
Browse files Browse the repository at this point in the history
The implementation follows the patterns established by the lowering of other
similar intrinsics.

In addition to the code for lowering, the DoTotalReduction template had to be
fixed to correctly break when signaled to do so by the accumulator function.

Differential Revision: https://reviews.llvm.org/D138140
  • Loading branch information
tarunprabhu committed Nov 29, 2022
1 parent 6841c43 commit d43a2f0
Show file tree
Hide file tree
Showing 5 changed files with 417 additions and 1 deletion.
13 changes: 13 additions & 0 deletions flang/include/flang/Optimizer/Builder/Runtime/Reduction.h
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,19 @@ mlir::Value genDotProduct(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Value vectorABox, mlir::Value vectorBBox,
mlir::Value resultBox);

/// Generate call to `Findloc` intrinsic runtime routine. This is the version
/// that does not take a dim argument.
void genFindloc(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value arrayBox, mlir::Value val,
mlir::Value maskBox, mlir::Value kind, mlir::Value back);

/// Generate call to `FindlocDim` intrinsic runtime routine. This is the version
/// that takes a dim argument.
void genFindlocDim(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value arrayBox, mlir::Value val,
mlir::Value dim, mlir::Value maskBox, mlir::Value kind,
mlir::Value back);

/// Generate call to `Maxloc` intrinsic runtime routine. This is the version
/// that does not take a dim argument.
void genMaxloc(fir::FirOpBuilder &builder, mlir::Location loc,
Expand Down
102 changes: 102 additions & 0 deletions flang/lib/Lower/IntrinsicCall.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -516,6 +516,7 @@ struct IntrinsicLibrary {
mlir::Value genIbits(mlir::Type, llvm::ArrayRef<mlir::Value>);
mlir::Value genIbset(mlir::Type, llvm::ArrayRef<mlir::Value>);
fir::ExtendedValue genIchar(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
fir::ExtendedValue genFindloc(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
mlir::Value genIeeeIsFinite(mlir::Type, llvm::ArrayRef<mlir::Value>);
template <mlir::arith::CmpIPredicate pred>
fir::ExtendedValue genIeeeTypeCompare(mlir::Type,
Expand Down Expand Up @@ -801,6 +802,15 @@ static constexpr IntrinsicHandler handlers[]{
{{{"status", asValue, handleDynamicOptional}}},
/*isElemental=*/false},
{"exponent", &I::genExponent},
{"findloc",
&I::genFindloc,
{{{"array", asBox},
{"value", asAddr},
{"dim", asValue},
{"mask", asBox, handleDynamicOptional},
{"kind", asValue},
{"back", asValue, handleDynamicOptional}}},
/*isElemental=*/false},
{"floor", &I::genFloor},
{"fraction", &I::genFraction},
{"get_command",
Expand Down Expand Up @@ -3181,6 +3191,98 @@ mlir::Value IntrinsicLibrary::genExponent(mlir::Type resultType,
fir::getBase(args[0])));
}

// FINDLOC
fir::ExtendedValue
IntrinsicLibrary::genFindloc(mlir::Type resultType,
llvm::ArrayRef<fir::ExtendedValue> args) {
assert(args.size() == 6);

llvm::StringRef errMsg = "unexpected result for Findloc";

// Handle required array argument
mlir::Value array = builder.createBox(loc, args[0]);
unsigned rank = fir::BoxValue(array).rank();
assert(rank >= 1);

// Handle required value argument
mlir::Value val = builder.createBox(loc, args[1]);

// Check if dim argument is present
bool absentDim = isStaticallyAbsent(args[2]);

// Handle optional mask argument
auto mask = isStaticallyAbsent(args[3])
? builder.create<fir::AbsentOp>(
loc, fir::BoxType::get(builder.getI1Type()))
: builder.createBox(loc, args[3]);

// Handle optional kind argument
auto kind = isStaticallyAbsent(args[4])
? builder.createIntegerConstant(
loc, builder.getIndexType(),
builder.getKindMap().defaultIntegerKind())
: fir::getBase(args[4]);

// Handle optional back argument
auto back = isStaticallyAbsent(args[5]) ? builder.createBool(loc, false)
: fir::getBase(args[5]);

if (!absentDim && rank == 1) {
// If dim argument is present and the array is rank 1, then the result is
// a scalar (since the the result is rank-1 or 0).
// Therefore, we use a scalar result descriptor with FindlocDim().
// Create mutable fir.box to be passed to the runtime for the result.
fir::MutableBoxValue resultMutableBox =
fir::factory::createTempMutableBox(builder, loc, resultType);
mlir::Value resultIrBox =
fir::factory::getMutableIRBox(builder, loc, resultMutableBox);
mlir::Value dim = fir::getBase(args[2]);

fir::runtime::genFindlocDim(builder, loc, resultIrBox, array, val, dim,
mask, kind, back);

// Handle cleanup of allocatable result descriptor and return
fir::ExtendedValue res =
fir::factory::genMutableBoxRead(builder, loc, resultMutableBox);
return res.match(
[&](const mlir::Value &addr) -> fir::ExtendedValue {
addCleanUpForTemp(loc, addr);
return builder.create<fir::LoadOp>(loc, resultType, addr);
},
[&](const auto &) -> fir::ExtendedValue {
fir::emitFatalError(loc, errMsg);
});
}

// The result will be an array. Create mutable fir.box to be passed to the
// runtime for the result.
mlir::Type resultArrayType =
builder.getVarLenSeqTy(resultType, absentDim ? 1 : rank - 1);
fir::MutableBoxValue resultMutableBox =
fir::factory::createTempMutableBox(builder, loc, resultArrayType);
mlir::Value resultIrBox =
fir::factory::getMutableIRBox(builder, loc, resultMutableBox);

if (absentDim) {
fir::runtime::genFindloc(builder, loc, resultIrBox, array, val, mask, kind,
back);
} else {
mlir::Value dim = fir::getBase(args[2]);
fir::runtime::genFindlocDim(builder, loc, resultIrBox, array, val, dim,
mask, kind, back);
}

return fir::factory::genMutableBoxRead(builder, loc, resultMutableBox)
.match(
[&](const fir::ArrayBoxValue &box) -> fir::ExtendedValue {
addCleanUpForTemp(loc, box.getAddr());
return box;
},
[&](const auto &) -> fir::ExtendedValue {
fir::emitFatalError(loc, errMsg);
});
}

// FLOOR
mlir::Value IntrinsicLibrary::genFloor(mlir::Type resultType,
llvm::ArrayRef<mlir::Value> args) {
Expand Down
35 changes: 35 additions & 0 deletions flang/lib/Optimizer/Builder/Runtime/Reduction.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -576,6 +576,41 @@ void fir::runtime::genCountDim(fir::FirOpBuilder &builder, mlir::Location loc,
builder.create<fir::CallOp>(loc, func, args);
}

/// Generate call to `Findloc` intrinsic runtime routine. This is the version
/// that does not take a dim argument.
void fir::runtime::genFindloc(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value arrayBox,
mlir::Value valBox, mlir::Value maskBox,
mlir::Value kind, mlir::Value back) {
auto func = fir::runtime::getRuntimeFunc<mkRTKey(Findloc)>(loc, builder);
auto fTy = func.getFunctionType();
auto sourceFile = fir::factory::locationToFilename(builder, loc);
auto sourceLine =
fir::factory::locationToLineNo(builder, loc, fTy.getInput(5));
auto args = fir::runtime::createArguments(builder, loc, fTy, resultBox,
arrayBox, valBox, kind, sourceFile,
sourceLine, maskBox, back);
builder.create<fir::CallOp>(loc, func, args);
}

/// Generate call to `FindlocDim` intrinsic runtime routine. This is the version
/// that takes a dim argument.
void fir::runtime::genFindlocDim(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value arrayBox,
mlir::Value valBox, mlir::Value dim,
mlir::Value maskBox, mlir::Value kind,
mlir::Value back) {
auto func = fir::runtime::getRuntimeFunc<mkRTKey(FindlocDim)>(loc, builder);
auto fTy = func.getFunctionType();
auto sourceFile = fir::factory::locationToFilename(builder, loc);
auto sourceLine =
fir::factory::locationToLineNo(builder, loc, fTy.getInput(6));
auto args = fir::runtime::createArguments(
builder, loc, fTy, resultBox, arrayBox, valBox, kind, dim, sourceFile,
sourceLine, maskBox, back);
builder.create<fir::CallOp>(loc, func, args);
}

/// Generate call to `Maxloc` intrinsic runtime routine. This is the version
/// that does not take a dim argument.
void fir::runtime::genMaxloc(fir::FirOpBuilder &builder, mlir::Location loc,
Expand Down
3 changes: 2 additions & 1 deletion flang/runtime/reduction-templates.h
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,8 @@ inline void DoTotalReduction(const Descriptor &x, int dim,
for (auto elements{x.Elements()}; elements--;
x.IncrementSubscripts(xAt), mask->IncrementSubscripts(maskAt)) {
if (IsLogicalElementTrue(*mask, maskAt)) {
accumulator.template AccumulateAt<TYPE>(xAt);
if (!accumulator.template AccumulateAt<TYPE>(xAt))
break;
}
}
return;
Expand Down
Loading

0 comments on commit d43a2f0

Please sign in to comment.