1#pragma once
2
3#include <c10/core/Backend.h>
4#include <c10/util/Exception.h>
5
6#include <ostream>
7
8namespace c10 {
9enum class Layout : int8_t {
10 Strided,
11 Sparse,
12 SparseCsr,
13 Mkldnn,
14 SparseCsc,
15 SparseBsr,
16 SparseBsc,
17 NumOptions
18};
19
20constexpr auto kStrided = Layout::Strided;
21constexpr auto kSparse = Layout::Sparse;
22constexpr auto kSparseCsr = Layout::SparseCsr;
23constexpr auto kMkldnn = Layout::Mkldnn;
24constexpr auto kSparseCsc = Layout::SparseCsc;
25constexpr auto kSparseBsr = Layout::SparseBsr;
26constexpr auto kSparseBsc = Layout::SparseBsc;
27
28inline 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
48inline 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