#pragma once
|
|
#include "caffe2/opt/backend_transformer_base.h"
|
|
#include <unordered_set>
|
|
C10_DECLARE_bool(caffe2_tvm_profiling_based_jit);
|
|
namespace caffe2 {
|
|
struct TvmTransformOptions final : public BackendTransformOptions {
|
explicit TvmTransformOptions() : BackendTransformOptions() {}
|
|
// Whether to enable profiling based jit
|
bool profiling_based_jit{true};
|
};
|
|
class CAFFE2_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 blacklisted_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 std::unordered_map<std::string, TensorShape>& shape_hints,
|
const std::unordered_set<int>& blacklisted_ops) override;
|
|
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 cluser connected TVM runnable ops into one TVMJitOp
|
NetDef applyTvmTransform(
|
NetDef* pred_net,
|
const std::unordered_set<std::string>& weights,
|
const std::unordered_set<int>& blacklisted_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.
|
CAFFE2_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 std::unordered_map<std::string, TensorShape>& shape_hints,
|
const std::unordered_set<int>& blacklisted_ops,
|
size_t max_batch_size,
|
size_t max_seq_size,
|
bool debug);
|
|
} // namespace caffe2
|
