1 | // Generated from "/code/pytorch/third_party/nvfuser/runtime/tensor.cu" |
2 | // 2023-02-12 08:01:26 |
3 | |
4 | namespace nvfuser_resources { |
5 | |
6 | constexpr const char* tensor_cu = R"( |
7 | template <typename T, int N> |
8 | struct Tensor { |
9 | __device__ T& operator[](nvfuser_index_t ind) { |
10 | return data[ind]; |
11 | }; |
12 | |
13 | T* data; |
14 | nvfuser_index_t size[N]; |
15 | nvfuser_index_t stride[N]; |
16 | }; |
17 | |
18 | // Specialization for 0-dim case as it does not need size and stride arrays. |
19 | // They will be an error as well since zero-length arrays are not allowed. |
20 | template <typename T> |
21 | struct Tensor<T, 0> { |
22 | __device__ T& operator[](nvfuser_index_t) { |
23 | return *data; |
24 | }; |
25 | |
26 | T* data; |
27 | }; |
28 | |
29 | // Specialization for 0-dim case that's easy to pass in a CPU based tensor. |
30 | template <typename T> |
31 | struct CpuScalarTensor { |
32 | __device__ T& operator[](int) { |
33 | return data; |
34 | }; |
35 | |
36 | T data; |
37 | }; |
38 | )" ; |
39 | |
40 | } // namespace nvfuser_resources |
41 | |