FM / opt

/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/opt/

/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/opt

Name	Size	Mode	Actions
annotations.h	2141	0644	edit dl rm
backend_cutting.h	586	0644	edit dl rm
backend_transformer_base.h	2773	0644	edit dl rm
bound_shape_inferencer.h	5884	0644	edit dl rm
converter.h	3013	0644	edit dl rm
device.h	442	0644	edit dl rm
distributed.h	1126	0644	edit dl rm
fakefp16_transform.h	584	0644	edit dl rm
fusion.h	4141	0644	edit dl rm
glow_net_transform.h	1625	0644	edit dl rm
mobile.h	388	0644	edit dl rm
onnxifi_op.h	19497	0644	edit dl rm
onnxifi_transformer.h	6859	0644	edit dl rm
onnx_convert.h	1279	0644	edit dl rm
optimizer.h	402	0644	edit dl rm
optimize_ideep.h	351	0644	edit dl rm
passes.h	2642	0644	edit dl rm
shape_info.h	4784	0644	edit dl rm
tvm_transformer.h	2992	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/opt/tvm_transformer.h (2992B)

#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

cmd: