31 lines
868 B
C++
31 lines
868 B
C++
#pragma once
|
|
|
|
#include "mlir/Pass/Pass.h"
|
|
|
|
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
|
|
#include "src/Compiler/CompilerOptions.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
namespace pim {
|
|
|
|
struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
|
|
StringRef getArgument() const override { return "bufferize-pim"; }
|
|
StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; }
|
|
|
|
PimBufferizationPass() = default;
|
|
PimBufferizationPass(const PimBufferizationPass& pass) {}
|
|
|
|
void runOnOperation() final;
|
|
|
|
private:
|
|
void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
|
|
};
|
|
|
|
} // namespace pim
|
|
|
|
std::unique_ptr<Pass> createBufferizePimPass() { return std::make_unique<pim::PimBufferizationPass>(); }
|
|
|
|
} // namespace onnx_mlir
|