1 | #pragma once |
2 | |
3 | #include <c10/core/Backend.h> |
4 | #include <c10/util/Exception.h> |
5 | |
6 | #include <ostream> |
7 | |
8 | namespace c10 { |
9 | enum class Layout : int8_t { |
10 | Strided, |
11 | Sparse, |
12 | SparseCsr, |
13 | Mkldnn, |
14 | SparseCsc, |
15 | SparseBsr, |
16 | SparseBsc, |
17 | NumOptions |
18 | }; |
19 | |
20 | constexpr auto kStrided = Layout::Strided; |
21 | constexpr auto kSparse = Layout::Sparse; |
22 | constexpr auto kSparseCsr = Layout::SparseCsr; |
23 | constexpr auto kMkldnn = Layout::Mkldnn; |
24 | constexpr auto kSparseCsc = Layout::SparseCsc; |
25 | constexpr auto kSparseBsr = Layout::SparseBsr; |
26 | constexpr auto kSparseBsc = Layout::SparseBsc; |
27 | |
28 | inline Layout layout_from_backend(Backend backend) { |
29 | switch (backend) { |
30 | case Backend::SparseCPU: |
31 | case Backend::SparseCUDA: |
32 | case Backend::SparseHIP: |
33 | case Backend::SparseVE: |
34 | case Backend::SparseXPU: |
35 | return Layout::Sparse; |
36 | case Backend::MkldnnCPU: |
37 | return Layout::Mkldnn; |
38 | case Backend::SparseCsrCPU: |
39 | case Backend::SparseCsrCUDA: |
40 | TORCH_CHECK( |
41 | false, |
42 | "Cannot map Backend SparseCsrCPU|SparseCsrCUDA to a unique layout." ); |
43 | default: |
44 | return Layout::Strided; |
45 | } |
46 | } |
47 | |
48 | inline std::ostream& operator<<(std::ostream& stream, at::Layout layout) { |
49 | switch (layout) { |
50 | case at::kStrided: |
51 | return stream << "Strided" ; |
52 | case at::kSparse: |
53 | return stream << "Sparse" ; |
54 | case at::kSparseCsr: |
55 | return stream << "SparseCsr" ; |
56 | case at::kSparseCsc: |
57 | return stream << "SparseCsc" ; |
58 | case at::kSparseBsr: |
59 | return stream << "SparseBsr" ; |
60 | case at::kSparseBsc: |
61 | return stream << "SparseBsc" ; |
62 | case at::kMkldnn: |
63 | return stream << "Mkldnn" ; |
64 | default: |
65 | TORCH_CHECK(false, "Unknown layout" ); |
66 | } |
67 | } |
68 | |
69 | } // namespace c10 |
70 | |