138 lines
4.6 KiB
C++
138 lines
4.6 KiB
C++
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Pass/PassManager.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "mlir/Transforms/Passes.h"
|
|
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/raw_os_ostream.h"
|
|
|
|
#include <fstream>
|
|
|
|
#include "Common/PimCommon.hpp"
|
|
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
|
|
#include "src/Compiler/CompilerOptions.hpp"
|
|
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
bool haveSameStaticShape(Value lhs, Value rhs);
|
|
|
|
namespace {
|
|
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
|
|
|
|
struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp>> {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToSpatialPass)
|
|
StringRef getArgument() const override { return "convert-onnx-to-spatial"; }
|
|
StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
|
|
|
|
ONNXToSpatialPass() = default;
|
|
ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
|
|
|
|
void runOnOperation() override;
|
|
|
|
private:
|
|
void annotateWeightsConstants(func::FuncOp funcOp) const;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void ONNXToSpatialPass::runOnOperation() {
|
|
ModuleOp moduleOp = getOperation();
|
|
MLIRContext* ctx = &getContext();
|
|
|
|
RewritePatternSet mergeActivationPatterns(ctx);
|
|
mergeActivationPatterns.add<onnxToArithConstantOp>(ctx);
|
|
mergeActivationPatterns.add<convAddToConvWithBiasPatternLeft>(ctx);
|
|
mergeActivationPatterns.add<convAddToConvWithBiasPatternRight>(ctx);
|
|
mergeActivationPatterns.add<matMulAddToGemmPattern>(ctx);
|
|
mergeActivationPatterns.add<matMulToGemmPattern>(ctx);
|
|
mergeActivationPatterns.add<removeFlattenSameShapePattern>(ctx);
|
|
populateMatMulRewritePatterns(mergeActivationPatterns, ctx);
|
|
|
|
if (failed(applyPatternsGreedily(moduleOp, std::move(mergeActivationPatterns))))
|
|
llvm::dbgs() << "Failed to merge activation patterns, continuing...\n";
|
|
|
|
IRRewriter rewriter(moduleOp);
|
|
auto entryFunc = getPimEntryFunc(moduleOp);
|
|
if (failed(entryFunc)) {
|
|
signalPassFailure();
|
|
return;
|
|
}
|
|
|
|
ConversionTarget target(*ctx);
|
|
target.addLegalDialect<spatial::SpatialDialect, ONNXDialect, tensor::TensorDialect, arith::ArithDialect>();
|
|
target.addDynamicallyLegalOp<ONNXMatMulOp>([](ONNXMatMulOp op) {
|
|
return cast<ShapedType>(op.getY().getType()).getRank() != 2;
|
|
});
|
|
target.addIllegalOp<ONNXGemmOp>();
|
|
target.addIllegalOp<ONNXConvOp>();
|
|
target.addIllegalOp<ONNXLRNOp>();
|
|
target.addIllegalOp<ONNXMaxPoolSingleOutOp>();
|
|
target.addIllegalOp<ONNXAveragePoolOp>();
|
|
target.addIllegalOp<ONNXConcatOp>();
|
|
target.addIllegalOp<ONNXSoftmaxOp>();
|
|
target.addIllegalOp<ONNXReduceMeanV13Op>();
|
|
target.addIllegalOp<ONNXReshapeOp>();
|
|
|
|
RewritePatternSet patterns(ctx);
|
|
patterns.add<removeLRNPattern>(ctx);
|
|
|
|
populateConvOpPatterns(patterns, ctx);
|
|
populatePoolTilingPattern(patterns, ctx);
|
|
populateOnnxGemmOpPatterns(patterns, ctx);
|
|
populateReshapeConversionPattern(patterns, ctx);
|
|
|
|
populateONNXConcatToTensorConcatPattern(patterns, ctx);
|
|
|
|
if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
|
|
signalPassFailure();
|
|
return;
|
|
}
|
|
|
|
// Count the number of compute ops and check they do not exceed the core count
|
|
if (coresCount != -1) {
|
|
int computeOpsCount = 0;
|
|
for (auto& op : entryFunc->getFunctionBody().front().getOperations())
|
|
if (isa<spatial::SpatWeightedCompute>(op))
|
|
computeOpsCount++;
|
|
|
|
if (computeOpsCount > coresCount) {
|
|
llvm::dbgs() << "Number of compute ops exceeds the core count\n";
|
|
signalPassFailure();
|
|
return;
|
|
}
|
|
}
|
|
|
|
PassManager cleanupPM(ctx);
|
|
cleanupPM.addPass(createCanonicalizerPass());
|
|
if (failed(cleanupPM.run(moduleOp)))
|
|
llvm::dbgs() << "Failed to run canonicalization cleanup, continuing...\n";
|
|
|
|
annotateWeightsConstants(*entryFunc);
|
|
|
|
// Dump to file for debug
|
|
dumpModule(moduleOp, "spatial");
|
|
}
|
|
|
|
void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
|
|
funcOp.walk([&](arith::ConstantOp constantOp) {
|
|
bool isAlwaysWeight =
|
|
llvm::all_of(constantOp->getUsers(), [](auto user) -> bool { return isa<spatial::SpatWeightedCompute>(user); });
|
|
if (isAlwaysWeight)
|
|
markWeightAlways(constantOp);
|
|
});
|
|
}
|
|
|
|
std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
|
|
|
|
} // namespace onnx_mlir
|