xref: /aosp_15_r20/external/pytorch/torch/csrc/utils/tensor_layouts.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <ATen/Layout.h>
2 #include <c10/core/ScalarType.h>
3 #include <torch/csrc/DynamicTypes.h>
4 #include <torch/csrc/Exceptions.h>
5 #include <torch/csrc/Layout.h>
6 #include <torch/csrc/python_headers.h>
7 #include <torch/csrc/utils/object_ptr.h>
8 #include <torch/csrc/utils/tensor_layouts.h>
9 
10 namespace torch::utils {
11 
12 #define REGISTER_LAYOUT(layout, LAYOUT)                                     \
13   PyObject* layout##_layout =                                               \
14       THPLayout_New(at::Layout::LAYOUT, "torch." #layout);                  \
15   Py_INCREF(layout##_layout);                                               \
16   if (PyModule_AddObject(torch_module, "" #layout, layout##_layout) != 0) { \
17     throw python_error();                                                   \
18   }                                                                         \
19   registerLayoutObject((THPLayout*)layout##_layout, at::Layout::LAYOUT);
20 
initializeLayouts()21 void initializeLayouts() {
22   auto torch_module = THPObjectPtr(PyImport_ImportModule("torch"));
23   if (!torch_module)
24     throw python_error();
25 
26   PyObject* strided_layout =
27       THPLayout_New(at::Layout::Strided, "torch.strided");
28   Py_INCREF(strided_layout);
29   if (PyModule_AddObject(torch_module, "strided", strided_layout) != 0) {
30     throw python_error();
31   }
32   registerLayoutObject((THPLayout*)strided_layout, at::Layout::Strided);
33 
34   PyObject* sparse_coo_layout =
35       THPLayout_New(at::Layout::Sparse, "torch.sparse_coo");
36   Py_INCREF(sparse_coo_layout);
37   if (PyModule_AddObject(torch_module, "sparse_coo", sparse_coo_layout) != 0) {
38     throw python_error();
39   }
40   registerLayoutObject((THPLayout*)sparse_coo_layout, at::Layout::Sparse);
41 
42   REGISTER_LAYOUT(sparse_csr, SparseCsr)
43   REGISTER_LAYOUT(sparse_csc, SparseCsc)
44   REGISTER_LAYOUT(sparse_bsr, SparseBsr)
45   REGISTER_LAYOUT(sparse_bsc, SparseBsc)
46 
47   PyObject* mkldnn_layout = THPLayout_New(at::Layout::Mkldnn, "torch._mkldnn");
48   Py_INCREF(mkldnn_layout);
49   if (PyModule_AddObject(torch_module, "_mkldnn", mkldnn_layout) != 0) {
50     throw python_error();
51   }
52   registerLayoutObject((THPLayout*)mkldnn_layout, at::Layout::Mkldnn);
53 
54   REGISTER_LAYOUT(jagged, Jagged);
55 }
56 
57 } // namespace torch::utils
58