35 lines
959 B
C++
35 lines
959 B
C++
#pragma once
|
|
|
|
#include "mlir/Pass/Pass.h"
|
|
|
|
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
using namespace mlir;
|
|
extern bool haveSameStaticShape(Value lhs, Value rhs);
|
|
|
|
namespace spatial {
|
|
|
|
#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 spatial
|
|
|
|
std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<spatial::ONNXToSpatialPass>(); }
|
|
|
|
} // namespace onnx_mlir
|