yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Sai Praveen BangaruFix issue with slang-embed & include ordering (#5680)e6cf93e3e

master
branchesmasterall refs →
5.0 KiB181 linesraw
1// Prelude for PyTorch cpp binding.
2
3// clang-format off
4#include <torch/extension.h>
5// clang-format on
6
7#include <ATen/cuda/CUDAContext.h>
8#include <ATen/cuda/CUDAUtils.h>
9#include <stdexcept>
10#include <string>
11#include <vector>
12
13#ifdef SLANG_LLVM
14#include "slang-llvm.h"
15#else // SLANG_LLVM
16#if SLANG_GCC_FAMILY && __GNUC__ < 6
17#include <cmath>
18#define SLANG_PRELUDE_STD std::
19#else
20#include <math.h>
21#define SLANG_PRELUDE_STD
22#endif
23
24#include <assert.h>
25#include <stdint.h>
26#include <stdlib.h>
27#include <string.h>
28#endif // SLANG_LLVM
29
30#include "../source/core/slang-string.h"
31
32#if defined(_MSC_VER)
33#define SLANG_PRELUDE_SHARED_LIB_EXPORT __declspec(dllexport)
34#else
35#define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__((__visibility__("default")))
36// #   define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__ ((dllexport))
37// __attribute__((__visibility__("default")))
38#endif
39
40#ifdef __cplusplus
41#define SLANG_PRELUDE_EXTERN_C extern "C"
42#define SLANG_PRELUDE_EXTERN_C_START \
43    extern "C"                       \
44    {
45#define SLANG_PRELUDE_EXTERN_C_END }
46#else
47#define SLANG_PRELUDE_EXTERN_C
48#define SLANG_PRELUDE_EXTERN_C_START
49#define SLANG_PRELUDE_EXTERN_C_END
50#endif
51
52#define SLANG_PRELUDE_NAMESPACE
53
54#ifndef SLANG_NO_THROW
55#define SLANG_NO_THROW
56#endif
57#ifndef SLANG_STDCALL
58#define SLANG_STDCALL
59#endif
60#ifndef SLANG_MCALL
61#define SLANG_MCALL SLANG_STDCALL
62#endif
63#ifndef SLANG_FORCE_INLINE
64#define SLANG_FORCE_INLINE inline
65#endif
66#include "slang-cpp-scalar-intrinsics.h"
67#include "slang-cpp-types-core.h"
68
69
70static const int kSlangTorchTensorMaxDim = 5;
71
72struct TensorView
73{
74    uint8_t* data;
75    uint32_t strides[kSlangTorchTensorMaxDim];
76    uint32_t sizes[kSlangTorchTensorMaxDim];
77    uint32_t dimensionCount;
78};
79
80
81TensorView make_tensor_view(
82    torch::Tensor val,
83    const char* name,
84    torch::ScalarType targetScalarType,
85    bool requireContiguous)
86{
87    // We're currently not trying to implicitly cast or transfer to device for two reasons:
88    // 1. There appears to be a bug with .to() where successive calls after the first one fail.
89    // 2. Silent casts like this can cause large memory allocations & unexpected overheads.
90    //    It's better to be explicit.
91
92    // Expect tensors to be on CUDA device
93    if (!val.device().is_cuda())
94        throw std::runtime_error(
95            std::string(name).append(": tensor is not on CUDA device.").c_str());
96
97    // Expect tensors to be the right type.
98    if (val.dtype() != targetScalarType)
99        throw std::runtime_error(
100            std::string(name).append(": tensor is not of the expected type.").c_str());
101
102    // Check that the tensor is contiguous
103    if (requireContiguous && !val.is_contiguous())
104        throw std::runtime_error(std::string(name).append(": tensor is not contiguous.").c_str());
105
106    TensorView res = {};
107    res.dimensionCount = val.dim();
108    res.data = nullptr;
109    size_t elementSize = 4;
110
111    switch (val.scalar_type())
112    {
113    case torch::kInt8:
114    case torch::kUInt8:
115        elementSize = 1;
116        res.data = (uint8_t*)val.data_ptr<uint8_t>();
117        break;
118    case torch::kBFloat16:
119        elementSize = 2;
120        res.data = (uint8_t*)val.data_ptr<torch::BFloat16>();
121        break;
122    case torch::kFloat16:
123        elementSize = 2;
124        res.data = (uint8_t*)val.data_ptr<at::Half>();
125        break;
126    case torch::kInt16:
127        elementSize = 2;
128        res.data = (uint8_t*)val.data_ptr<int16_t>();
129        break;
130    case torch::kFloat32:
131        elementSize = 4;
132        res.data = (uint8_t*)val.data_ptr<float>();
133        break;
134    case torch::kInt32:
135        elementSize = 4;
136        res.data = (uint8_t*)val.data_ptr<int32_t>();
137        break;
138    case torch::kFloat64:
139        elementSize = 8;
140        res.data = (uint8_t*)val.data_ptr<double>();
141        break;
142    case torch::kInt64:
143        elementSize = 8;
144        res.data = (uint8_t*)val.data_ptr<int64_t>();
145        break;
146    case torch::kBool:
147        elementSize = 1;
148        res.data = (uint8_t*)val.data_ptr<bool>();
149        break;
150    }
151
152    if (val.dim() > kSlangTorchTensorMaxDim)
153        throw std::runtime_error(std::string(name)
154                                     .append(": number of dimensions exceeds limit (")
155                                     .append(std::to_string(kSlangTorchTensorMaxDim))
156                                     .append(")")
157                                     .c_str());
158
159    bool isEmpty = true;
160    for (int i = 0; i < val.dim(); ++i)
161    {
162        res.strides[i] = val.stride(i) * elementSize;
163        if (res.strides[i] == 0)
164            throw std::runtime_error(
165                std::string(name)
166                    .append(": tensors with broadcasted dimensions are not supported (use "
167                            "tensor.contiguous() to make tensor whole)")
168                    .c_str());
169
170        res.sizes[i] = val.size(i);
171        if (res.sizes[i] > 0)
172            isEmpty = false;
173    }
174
175    if (!res.data && !isEmpty)
176        throw std::runtime_error(std::string(name).append(": data pointer is invalid.").c_str());
177
178    return res;
179}
180
181#define SLANG_PRELUDE_EXPORT