94 lines
2.9 KiB
C++
94 lines
2.9 KiB
C++
#pragma once
|
|
|
|
#include "caffe2/opt/backend_transformer_base.h"
|
|
|
|
#include <unordered_set>
|
|
|
|
namespace caffe2 {
|
|
|
|
struct TvmTransformOptions final : public BackendTransformOptions {
|
|
explicit TvmTransformOptions() : BackendTransformOptions() {}
|
|
|
|
// Whether to enable profiling based jit
|
|
bool profiling_based_jit{false};
|
|
};
|
|
|
|
class TORCH_API TvmTransformer final : public BackendTransformerBase {
|
|
public:
|
|
explicit TvmTransformer(const TvmTransformOptions& opts)
|
|
: BackendTransformerBase(), opts_(opts) {}
|
|
|
|
~TvmTransformer() override {}
|
|
|
|
// Given workspace and predict net, cluster continuous parts that can be run
|
|
// by TVM and create one TVMJit op for each clustered subgraph.
|
|
// \param ws c2 workspace
|
|
// \param pred_net c2 predict net
|
|
// \param weight_names list of the names of the constant weights
|
|
// \param shape_hints User provided shape info, usually for primary inputs so
|
|
// that bound shape inference can have something to start
|
|
// \param blocklisted_ops a set of ops that we don't want to lower to TVM in
|
|
// terms of their net positions. This is very useful for debugging but for
|
|
// normal runs it should be empty
|
|
void transform(
|
|
Workspace* ws,
|
|
NetDef* pred_net,
|
|
const std::vector<std::string>& weight_names,
|
|
const ShapeInfoMap& shape_hints,
|
|
const std::unordered_set<int>& blocklisted_ops) override;
|
|
|
|
static const std::unordered_set<std::string>& getSupportedOps();
|
|
|
|
static bool canConvertFullGraph(
|
|
const caffe2::NetDef& net,
|
|
const std::unordered_set<int>& blocklisted_ops);
|
|
|
|
private:
|
|
// Given TVM runnable subnets, contract them into one TVMJitOp
|
|
NetDef buildTvmOp(
|
|
const caffe2::NetDef& net,
|
|
const std::unordered_set<std::string>& weights,
|
|
const ShapeInfoMap& shape_hints);
|
|
|
|
// Apply transform to cluster connected TVM runnable ops into one TVMJitOp
|
|
NetDef applyTvmTransform(
|
|
NetDef* pred_net,
|
|
const std::unordered_set<std::string>& weights,
|
|
const std::unordered_set<int>& blocklisted_ops,
|
|
const ShapeInfoMap& shape_hints);
|
|
|
|
// Options
|
|
TvmTransformOptions opts_;
|
|
|
|
// Track number of TVMJitOp we created
|
|
int tvm_op_id_{0};
|
|
|
|
// Model id
|
|
std::string model_id_;
|
|
};
|
|
|
|
// Helper function to clean up a net and run tvm transform.
|
|
TORCH_API void tvmTransform(
|
|
NetDef* net,
|
|
Workspace* ws,
|
|
const std::vector<std::string>& input_names,
|
|
const std::vector<std::string>& output_names,
|
|
const std::vector<std::string>& weight_names,
|
|
const ShapeInfoMap& shape_hints,
|
|
const std::unordered_set<int>& blocklisted_ops,
|
|
int32_t max_batch_size,
|
|
int32_t max_seq_size,
|
|
int32_t num_embeddings,
|
|
int32_t embedding_size,
|
|
int32_t tvm_min_ops,
|
|
bool tvm_profiling_based_jit,
|
|
bool debug);
|
|
|
|
TORCH_API void cleanUpPredictNet(
|
|
NetDef* net,
|
|
const std::vector<std::string>& input_names,
|
|
const std::vector<std::string>& output_names,
|
|
const std::vector<std::string>& weight_names);
|
|
|
|
} // namespace caffe2
|