NiccoloN
636310d0cb
Validate Operations / validate-operations (push) Has been cancelled
add shared checked arithmetic helpers refactor pim passes into Pim/Transforms more robust memory coalescing pass
97 lines
3.2 KiB
C++
97 lines
3.2 KiB
C++
#include "mlir/Dialect/SCF/IR/SCF.h"
|
|
|
|
#include "llvm/Support/MathExtras.h"
|
|
|
|
#include <optional>
|
|
|
|
#include "ConstantUtils.hpp"
|
|
#include "LoopUtils.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
namespace {
|
|
|
|
static std::optional<int64_t> getStaticTripCount(Value lowerBound, Value upperBound, Value step) {
|
|
auto lower = matchConstantIndexValue(lowerBound);
|
|
auto upper = matchConstantIndexValue(upperBound);
|
|
auto stepValue = matchConstantIndexValue(step);
|
|
if (!lower || !upper || !stepValue)
|
|
return std::nullopt;
|
|
if (*stepValue <= 0)
|
|
return std::nullopt;
|
|
if (*upper <= *lower)
|
|
return int64_t {0};
|
|
return llvm::divideCeil(*upper - *lower, *stepValue);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
static LogicalResult validateNormalizedLoopYields(Location loc, ValueRange initArgs, ArrayRef<Value> yieldedValues) {
|
|
if (yieldedValues.size() == initArgs.size())
|
|
return success();
|
|
|
|
emitError(loc) << "normalized loop body yielded " << yieldedValues.size() << " values for " << initArgs.size()
|
|
<< " iter args";
|
|
return failure();
|
|
}
|
|
|
|
FailureOr<NormalizedLoopResult> buildNormalizedScfFor(OpBuilder& builder,
|
|
Location loc,
|
|
Value lowerBound,
|
|
Value upperBound,
|
|
Value step,
|
|
ValueRange initArgs,
|
|
NormalizedLoopBodyBuilder bodyBuilder) {
|
|
NormalizedLoopResult result;
|
|
|
|
if (auto stepValue = matchConstantIndexValue(step); stepValue && *stepValue <= 0) {
|
|
emitError(loc) << "normalized scf.for requires a positive step, got " << *stepValue;
|
|
return failure();
|
|
}
|
|
|
|
if (auto tripCount = getStaticTripCount(lowerBound, upperBound, step)) {
|
|
if (*tripCount == 0) {
|
|
llvm::append_range(result.results, initArgs);
|
|
return result;
|
|
}
|
|
|
|
if (*tripCount == 1) {
|
|
result.inductionVar = lowerBound;
|
|
if (failed(bodyBuilder(builder, loc, lowerBound, initArgs, result.results)))
|
|
return failure();
|
|
if (failed(validateNormalizedLoopYields(loc, initArgs, result.results)))
|
|
return failure();
|
|
return result;
|
|
}
|
|
}
|
|
|
|
result.loop = scf::ForOp::create(builder, loc, lowerBound, upperBound, step, initArgs);
|
|
result.inductionVar = result.loop.getInductionVar();
|
|
|
|
{
|
|
OpBuilder::InsertionGuard guard(builder);
|
|
Block* body = result.loop.getBody();
|
|
if (!body->empty())
|
|
if (auto yieldOp = dyn_cast<scf::YieldOp>(body->back()))
|
|
yieldOp->erase();
|
|
builder.setInsertionPointToEnd(body);
|
|
ValueRange iterArgs = result.loop.getRegionIterArgs();
|
|
if (failed(bodyBuilder(builder, loc, result.inductionVar, iterArgs, result.results))) {
|
|
result.loop.erase();
|
|
return failure();
|
|
}
|
|
if (failed(validateNormalizedLoopYields(loc, initArgs, result.results))) {
|
|
result.loop.erase();
|
|
return failure();
|
|
}
|
|
scf::YieldOp::create(builder, loc, result.results);
|
|
}
|
|
builder.setInsertionPointAfter(result.loop);
|
|
result.results.assign(result.loop.getResults().begin(), result.loop.getResults().end());
|
|
return result;
|
|
}
|
|
|
|
} // namespace onnx_mlir
|