Implement custom Shape struct with __device__ support

This commit is contained in:
2025-11-23 19:21:06 +01:00
parent 82a0e7c19d
commit 1102aef293

View File

@@ -1,11 +1,71 @@
#pragma once
#ifndef __host__
#define __host__
#endif
#ifndef __device__
#define __device__
#endif
#include <format>
#include <vector>
namespace CUDANet {
typedef std::vector<size_t> Shape;
struct Shape {
static constexpr size_t MAX_DIMS = 8;
size_t dims[MAX_DIMS];
size_t ndim;
__host__ __device__ Shape() : ndim(0) {
for (int i = 0; i < MAX_DIMS; i++) dims[i] = 0;
}
__host__ Shape(std::initializer_list<size_t> list) : ndim(list.size()) {
if (ndim > MAX_DIMS) {
throw std::runtime_error("Too many dimensions");
}
size_t i = 0;
for (auto val : list) {
dims[i++] = val;
}
for (; i < MAX_DIMS; i++) dims[i] = 0;
}
__host__ Shape(const std::vector<size_t>& vec) : ndim(vec.size()) {
if (ndim > MAX_DIMS) {
throw std::runtime_error("Too many dimensions");
}
for (size_t i = 0; i < ndim; i++) {
dims[i] = vec[i];
}
for (size_t i = ndim; i < MAX_DIMS; i++) dims[i] = 0;
}
__host__ __device__ size_t operator[](size_t idx) const {
return dims[idx];
}
__host__ __device__ size_t& operator[](size_t idx) {
return dims[idx];
}
__host__ __device__ size_t size() const { return ndim; }
__host__ bool operator==(const Shape& other) const {
if (ndim != other.ndim) return false;
for (size_t i = 0; i < ndim; i++) {
if (dims[i] != other.dims[i]) return false;
}
return true;
}
__host__ bool operator!=(const Shape& other) const {
return !(*this == other);
}
};
std::string format_shape(const Shape& shape) {
std::string result;