1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #pragma once
10
11 #include <ostream>
12
13 namespace vkcompute {
14
15 // Convenience constexpr to attach semantic names to WHCN dimension index
16 namespace WHCN {
17
18 constexpr int32_t kWidthDim = 0;
19 constexpr int32_t kHeightDim = 1;
20 constexpr int32_t kChannelsDim = 2;
21
22 } // namespace WHCN
23
24 namespace utils {
25
26 //
27 // GPU Storage Options
28 //
29
30 /**
31 * The enum below is used to describe what type of GPU memory will be used to
32 * store a particular tensor's data.
33 *
34 * BUFFER means that a SSBO (Shader Storage Buffer Object) will be used.
35 * TEXTURE_3D means that a 3-dimensional image texture will be used.
36 * TEXTURE_2D means that a 2-dimensional image texture will be used.
37 *
38 * UNKNOWN is not expected to be used.
39 */
40 enum class StorageType : uint8_t {
41 BUFFER,
42 TEXTURE_3D,
43 TEXTURE_2D,
44 };
45
46 static constexpr StorageType kBuffer = StorageType::BUFFER;
47 static constexpr StorageType kTexture3D = StorageType::TEXTURE_3D;
48 static constexpr StorageType kTexture2D = StorageType::TEXTURE_2D;
49
50 /*
51 * A tensor's memory layout is defined in one of two ways:
52 *
53 * 1. If it's a buffer backed tensor, the memory layout is defined by its
54 * `dim_order`, and by extension its `strides`.
55 * 2. If it's a texture backed tensor, the memory layout is defined by the
56 * combination of its `axis_map` and its `packed_dim`.
57 *
58 * Providing explicit memory layout metadata upon tensor construction is not
59 * very convenient from an API perspective, so the `GPUMemoryLayout` serves as
60 * an abstraction that is used to determine how to initialize a tensor's layout
61 * metadata based on the developer's intent. A `GPUMemoryLayout` is provided to
62 * the constructor of `vTensor`, which will use it to determine how to set its
63 * `dim_order` if it's a buffer backed tensor, or how to set its `axis_map` and
64 * `packed_dim` if it's a texture backed tensor.
65 *
66 * Note that GPUMemoryLayout is not stored as a tensor property, as it does not
67 * have any meaning after the vTensor is constructed. After construction,
68 * methods such as `virtual_transpose()` may be used to modify the tensor's
69 * layout metadata that cannot be represented by any `GPUMemoryLayout` entry.
70 * Nonetheless, a "best guess" of the closest memory layout can be produced via
71 * the `estimate_memory_layout()` API of `vTensor`.
72 *
73 * Currently, only 3 memory layouts are provided, but more will be added in the
74 * future that will enable different functionality such as minimizing texture
75 * memory footprint.
76 */
77 enum class GPUMemoryLayout : uint8_t {
78 /*
79 * The below memory layouts will produce a `vTensor` with the following
80 * properties:
81 *
82 * 1. For buffer backed tensors, the `dim_order` will be the same as a
83 * contiguous dim order, but with the specified dim last in the dim order.
84 * 2. For texture backed tensors, the packed dim will be the specified dim.
85 * The axis map will be `{0, 1, 2, 2}`.
86 */
87 TENSOR_WIDTH_PACKED = 0u,
88 TENSOR_HEIGHT_PACKED = 1u,
89 TENSOR_CHANNELS_PACKED = 2u,
90 };
91
92 static constexpr GPUMemoryLayout kWidthPacked =
93 GPUMemoryLayout::TENSOR_WIDTH_PACKED;
94
95 static constexpr GPUMemoryLayout kHeightPacked =
96 GPUMemoryLayout::TENSOR_HEIGHT_PACKED;
97
98 static constexpr GPUMemoryLayout kChannelsPacked =
99 GPUMemoryLayout::TENSOR_CHANNELS_PACKED;
100
101 template <typename T>
to_packed_dim(const GPUMemoryLayout layout)102 T to_packed_dim(const GPUMemoryLayout layout) {
103 switch (layout) {
104 case kWidthPacked:
105 return 0;
106 case kHeightPacked:
107 return 1;
108 case kChannelsPacked:
109 return 2;
110 };
111 // Should be unreachable
112 return 0;
113 }
114
115 inline std::ostream& operator<<(
116 std::ostream& os,
117 const GPUMemoryLayout layout) {
118 switch (layout) {
119 case kWidthPacked:
120 os << "TENSOR_WIDTH_PACKED";
121 break;
122 case kHeightPacked:
123 os << "TENSOR_HEIGHT_PACKED";
124 break;
125 case kChannelsPacked:
126 os << "TENSOR_CHANNELS_PACKED";
127 break;
128 }
129 return os;
130 }
131
132 } // namespace utils
133 } // namespace vkcompute
134