mirror of
https://github.com/intel/llvm.git
synced 2026-01-27 23:13:59 +08:00
308571074c
This is an artifact from merging MLIR into LLVM, the file headers are now aligned with the rest of the project.
121 lines
4.3 KiB
C++
121 lines
4.3 KiB
C++
//===- AddDefaultStatsTestPass.cpp - Testing pass to add default stats ----===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file defines a testing pass to add default statistics nodes to every
|
|
// quantization eligible op. Useful for unit testing.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/QuantOps/QuantOps.h"
|
|
#include "mlir/Dialect/QuantOps/QuantTypes.h"
|
|
#include "mlir/IR/Attributes.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/Quantizer/Configurations/FxpMathConfig.h"
|
|
#include "mlir/Quantizer/Support/Configuration.h"
|
|
#include "mlir/Quantizer/Support/ConstraintAnalysisGraph.h"
|
|
#include "mlir/Quantizer/Support/ConstraintAnalysisGraphTraits.h"
|
|
#include "mlir/Quantizer/Transforms/Passes.h"
|
|
#include "mlir/Support/LogicalResult.h"
|
|
#include "llvm/Support/GraphWriter.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::quantizer;
|
|
using namespace mlir::quant;
|
|
|
|
namespace {
|
|
|
|
class AddDefaultStatsPass : public FunctionPass<AddDefaultStatsPass> {
|
|
public:
|
|
AddDefaultStatsPass() = default;
|
|
AddDefaultStatsPass(SolverContext &solverContext,
|
|
const TargetConfiguration &config)
|
|
: explicitSolverContext(&solverContext), explicitConfig(&config) {}
|
|
|
|
void runOnFunction() override;
|
|
void runWithConfig(SolverContext &solverContext,
|
|
const TargetConfiguration &config);
|
|
|
|
private:
|
|
SolverContext *explicitSolverContext = nullptr;
|
|
const TargetConfiguration *explicitConfig = nullptr;
|
|
};
|
|
|
|
} // end anonymous namespace
|
|
|
|
void AddDefaultStatsPass::runOnFunction() {
|
|
if (explicitSolverContext && explicitConfig) {
|
|
// If explicitly constructed with a config and context.
|
|
runWithConfig(*explicitSolverContext, *explicitConfig);
|
|
return;
|
|
}
|
|
// For global pass registration, use defaults.
|
|
SolverContext solverContext(*getFunction().getContext());
|
|
auto config = FxpMathTargetConfig::create(solverContext);
|
|
runWithConfig(solverContext, *config);
|
|
}
|
|
|
|
void AddDefaultStatsPass::runWithConfig(SolverContext &solverContext,
|
|
const TargetConfiguration &config) {
|
|
auto func = getFunction();
|
|
|
|
// Insert stats for each argument.
|
|
for (auto arg : func.getArguments()) {
|
|
if (!config.isHandledType(arg.getType()))
|
|
continue;
|
|
OpBuilder b(func.getBody());
|
|
APFloat minValue(-1.0f);
|
|
APFloat maxValue(1.0f);
|
|
ElementsAttr layerStats = DenseFPElementsAttr::get(
|
|
RankedTensorType::get({2}, b.getF32Type()), {minValue, maxValue});
|
|
auto statsOp = b.create<StatisticsOp>(func.getLoc(), arg, layerStats,
|
|
nullptr, nullptr);
|
|
arg.replaceAllUsesWith(statsOp);
|
|
|
|
// StatsOp contained a use to 'arg' so make sure to reset it after replacing
|
|
// all of the uses of 'arg'.
|
|
statsOp.getOperation()->replaceUsesOfWith(statsOp, arg);
|
|
}
|
|
|
|
// Walk the ops and insert stats.
|
|
func.walk([&](Operation *op) {
|
|
if (!config.isRequireStatsOp(op)) {
|
|
return;
|
|
}
|
|
assert(op->getNumResults() == 1);
|
|
|
|
auto originalResult = op->getResult(0);
|
|
if (!config.isHandledType(originalResult.getType()))
|
|
return;
|
|
|
|
OpBuilder b(op->getBlock(), ++op->getIterator());
|
|
|
|
APFloat minValue(-1.0f);
|
|
APFloat maxValue(1.0f);
|
|
ElementsAttr layerStats = DenseFPElementsAttr::get(
|
|
RankedTensorType::get({2}, b.getF32Type()), {minValue, maxValue});
|
|
auto statsOp = b.create<StatisticsOp>(op->getLoc(), op->getResult(0),
|
|
layerStats, nullptr, nullptr);
|
|
originalResult.replaceAllUsesWith(statsOp);
|
|
|
|
// StatsOp contained a use to 'op' so make sure to reset it after replacing
|
|
// all of the uses of 'op'.
|
|
statsOp.getOperation()->replaceUsesOfWith(statsOp, originalResult);
|
|
});
|
|
}
|
|
|
|
std::unique_ptr<OpPassBase<FuncOp>>
|
|
mlir::quantizer::createAddDefaultStatsPass() {
|
|
return std::make_unique<AddDefaultStatsPass>();
|
|
}
|
|
|
|
static PassRegistration<AddDefaultStatsPass> pass(
|
|
"quantizer-add-default-stats-test",
|
|
"Adds default (dummy) statistics to all ops that can benefit from "
|
|
"runtime statistics. This is meant to help in early stage bootstrapping.");
|