1#pragma once
2
3#include <cstdint>
4#include <vector>
5
6#include "taichi/inc/constants.h"
7#include "taichi/ir/type_utils.h"
8#include "taichi/rhi/device.h"
9
10namespace taichi::lang {
11
12class Program;
13class Ndarray;
14class SNode;
15
16std::pair<DataType, uint32_t> buffer_format2type_channels(BufferFormat format);
17BufferFormat type_channels2buffer_format(const DataType &type,
18 uint32_t num_channels);
19
20class TI_DLL_EXPORT Texture {
21 public:
22 /* Constructs a Texture managed by Program.
23 * Texture object allocation and deallocation is handled by Program.
24 */
25 explicit Texture(Program *prog,
26 BufferFormat format,
27 int width,
28 int height,
29 int depth = 1);
30
31 /* Constructs a Texture from an existing DeviceAllocation
32 * It doesn't handle the allocation and deallocation.
33 */
34 explicit Texture(DeviceAllocation &devalloc,
35 BufferFormat format,
36 int width,
37 int height,
38 int depth = 1);
39
40 intptr_t get_device_allocation_ptr_as_int() const;
41
42 void from_ndarray(Ndarray *ndarray);
43
44 void from_snode(SNode *snode);
45
46 DeviceAllocation get_device_allocation() const {
47 return texture_alloc_;
48 }
49
50 ~Texture();
51
52 BufferFormat get_buffer_format() const {
53 return format_;
54 }
55
56 std::array<int, 3> get_size() const {
57 return {width_, height_, depth_};
58 }
59
60 private:
61 DeviceAllocation texture_alloc_{kDeviceNullAllocation};
62 DataType dtype_;
63 BufferFormat format_;
64 int num_channels_{4};
65 int width_;
66 int height_;
67 int depth_;
68
69 Program *prog_{nullptr};
70};
71
72} // namespace taichi::lang
73