File size: 3,540 Bytes
e873e70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | #include <sycl/sycl.hpp>
#include <torch/torch.h>
#include <cstdint>
using namespace sycl;
// Multi-dtype XPU (SYCL) ReLU: float32 / float16 / bfloat16 / int8.
// Vectorized: each work-item processes a 16-byte (128-bit) chunk -> coalesced,
// dword-aligned access. On the Intel iGPU this reaches the memory wall for every
// dtype; the naive 1-element/work-item kernel left sub-dword types issue-bound
// (int8 ~2.6x slower, fp16 ~1.2x). fp16 & bf16 are 16-bit sign-magnitude floats,
// so relu = "zero the lane if the sign bit is set" on the raw uint16.
namespace {
struct alignas(16) V16 { uint32_t w[4]; }; // 16-byte vector
template <typename T> inline T relu_dev(T x);
template <> inline float relu_dev<float>(float x) { return x > 0.f ? x : 0.f; }
template <> inline uint16_t relu_dev<uint16_t>(uint16_t x) { return (x & 0x8000u) ? uint16_t(0) : x; }
template <> inline int8_t relu_dev<int8_t>(int8_t x) { return x > 0 ? x : int8_t(0); }
template <typename T>
void launch(sycl::queue& q, T* o, const T* in, int64_t n) {
constexpr int C = 16 / sizeof(T); // elements per 16-byte chunk
const int64_t nv = n / C;
const int64_t tail_start = nv * C;
if (nv > 0) {
auto* vin = reinterpret_cast<const V16*>(in);
auto* vout = reinterpret_cast<V16*>(o);
q.parallel_for(range<1>(nv), [=](id<1> idx) {
int64_t i = idx[0];
V16 raw = vin[i];
T* e = reinterpret_cast<T*>(&raw);
#pragma unroll
for (int k = 0; k < C; ++k) e[k] = relu_dev<T>(e[k]);
vout[i] = raw;
});
}
if (tail_start < n) {
const int64_t rem = n - tail_start;
q.parallel_for(range<1>(rem), [=](id<1> idx) {
int64_t i = tail_start + idx[0];
o[i] = relu_dev<T>(in[i]);
});
}
q.wait();
}
} // namespace
void relu(torch::Tensor& out, const torch::Tensor& input) {
TORCH_CHECK(input.device().is_xpu(), "input must be a XPU tensor");
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
TORCH_CHECK(out.is_contiguous(), "output must be contiguous");
TORCH_CHECK(input.sizes() == out.sizes(),
"Tensors must have the same shape. Got input shape: ",
input.sizes(), " and output shape: ", out.sizes());
TORCH_CHECK(input.scalar_type() == out.scalar_type(),
"Tensors must have the same data type. Got input dtype: ",
input.scalar_type(), " and output dtype: ", out.scalar_type());
TORCH_CHECK(input.device() == out.device(),
"Tensors must be on the same device. Got input device: ",
input.device(), " and output device: ", out.device());
sycl::queue queue;
const int64_t n = input.numel();
switch (input.scalar_type()) {
case torch::kFloat:
launch<float>(queue, out.data_ptr<float>(), input.data_ptr<float>(), n);
break;
case torch::kHalf:
launch<uint16_t>(queue, reinterpret_cast<uint16_t*>(out.data_ptr<at::Half>()),
reinterpret_cast<const uint16_t*>(input.data_ptr<at::Half>()), n);
break;
case torch::kBFloat16:
launch<uint16_t>(queue, reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>()),
reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>()), n);
break;
case torch::kChar:
launch<int8_t>(queue, out.data_ptr<int8_t>(), input.data_ptr<int8_t>(), n);
break;
default:
TORCH_CHECK(false, "relu_xpu: unsupported dtype ", input.scalar_type(),
" (supported: float32, float16, bfloat16, int8)");
}
}
|