xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/codegen/fuser/tensor_info.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 #include <torch/csrc/Export.h>
3 
4 #include <cstddef>
5 #include <cstdint>
6 
7 namespace torch::jit::fuser {
8 
9 // Host-side view of TensorInfo
10 // Note dims[0] - we need to dynamically allocate the dims.
11 struct TORCH_API TensorInfo {
sizesTensorInfo12   uint32_t* sizes(size_t nDim) {
13     return &sizes_strides[0];
14   }
stridesTensorInfo15   uint32_t* strides(size_t nDim) {
16     return &sizes_strides[nDim];
17   }
18 
19   void* data;
20   // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
21   uint32_t sizes_strides[0];
22 };
23 
24 } // namespace torch::jit::fuser
25