Caffe2 - C++ API
A deep learning, cross platform ML framework
fused_kernel.h
1 #pragma once
2 
3 #include <ATen/ATen.h>
4 #include <torch/csrc/WindowsTorchApiMacro.h>
5 #include <torch/csrc/jit/fuser/cpu/dynamic_library.h>
6 #include <torch/csrc/jit/fuser/fused_kernel.h>
7 #include <torch/csrc/utils/disallow_copy.h>
8 
9 #include <cstdint>
10 #include <memory>
11 #include <string>
12 
13 namespace torch {
14 namespace jit {
15 namespace fuser {
16 namespace cpu {
17 
18 // Represents a compiled CPU kernel and the metadata necessary to run it
21  std::string name,
22  std::string code,
23  std::vector<TensorDesc> input_desc,
24  std::vector<TensorDesc> output_desc,
25  std::vector<PartitionDesc> chunk_desc,
26  std::vector<PartitionDesc> concat_desc,
27  bool has_random);
28 
29  at::Backend backend() const override {
30  return at::Backend::CPU;
31  }
32 
33  void launch_raw(const uint32_t numel, std::vector<void*>& arguments)
34  const override {
35  kernel(numel, arguments.data());
36  }
37 
38  private:
39  std::unique_ptr<DynamicLibrary> so_lib;
40  void (*kernel)(uint32_t, void**) = nullptr;
41 };
42 
43 } // namespace cpu
44 } // namespace fuser
45 } // namespace jit
46 } // namespace torch
Backend
This legacy enum class defines the set of backends supported by old school, code generated Type-based...
Definition: Backend.h:23
Definition: jit_type.h:17