WGPUBuffer buffer;

WGPUBufferUsage usage;

size_t size; // in bytes

static constexpr size_t kMaxRank = 8; // Maximum rank of a tensor, avoids

// dynamic allocation for shape data

std::array<size_t, kMaxRank> data = {0};

size_t rank = 0;

inline Shape() = default;

inline Shape(std::initializer_list<size_t> dims) {

assert(dims.size() <= kMaxRank);

std::copy(dims.begin(), dims.end(), data.begin());

rank = dims.size();

}

inline size_t &operator[](size_t index) {

assert(index < rank);

return data[index];

}

inline const size_t &operator[](size_t index) const {

assert(index < rank);

return data[index];

}

size_t numels = 1;

for (size_t i = 0; i < shape.rank; i++) {

numels *= shape.data[i];

}

return numels;

Tensor data; // non-owning view

size_t offset = 0;

size_t span = 0;

std::array<Tensor, N> data;

std::array<size_t, N> viewOffsets;

std::array<size_t, N> viewSpans;

Bindings(const std::initializer_list<Tensor> &init) {

std::copy(begin(init), end(init), begin(data));

std::fill(begin(viewOffsets), end(viewOffsets), 0);

for (size_t i = 0; i < N; ++i) {

viewSpans[i] = data[i].data.size;

}

}

Bindings(const std::array<Tensor, N> &init) {

std::copy(begin(init), end(init), begin(data));

std::fill(begin(viewOffsets), end(viewOffsets), 0);

for (size_t i = 0; i < N; ++i) {

viewSpans[i] = data[i].data.size;

}

}

Bindings(const std::initializer_list<TensorView> &init) {

size_t i = 0;

for (const auto &tv : init) {

data[i] = tv.data;

viewOffsets[i] = tv.offset;

viewSpans[i] = tv.span;

++i;

}

}

Bindings(const std::initializer_list<Array> &init) {

std::copy(begin(init), end(init), begin(data));

std::fill(begin(viewOffsets), end(viewOffsets), 0);

for (size_t i = 0; i < N; ++i) {

viewSpans[i] = data[i].size;

}

}

Tensor &operator[](std::size_t index) { return data[index]; }

const Tensor &operator[](std::size_t index) const { return data[index]; }

inline TensorPool(Context *ctx) : ctx(ctx), data() {};

Context *ctx;

std::unordered_map<WGPUBuffer, Tensor> data;

~TensorPool();

kf16, // (experimental)

kf32,

ki32

switch (type) {

case kf16:

return sizeof(uint16_t);

case kf32:

return sizeof(float);

case ki32:

return sizeof(int32_t);

default:

LOG(kDefLog, kError, "Invalid NumType in size calculation.");

return 0;

}

switch (type) {

case kf16:

return "f16";

case kf32:

return "f32";

case ki32:

return "i32";

default:

LOG(kDefLog, kError, "Invalid NumType in string conversion.");

return "unknown";

}

std::string str;

for (size_t i = 0; i < shape.rank; i++) {

str += std::to_string(shape.data[i]);

if (i < shape.rank - 1) {

str += ", ";

}

}

return str;

const std::string &to) {

size_t start_pos = 0;

while ((start_pos = str.find(from, start_pos)) != std::string::npos) {

str.replace(start_pos, from.length(), to);

start_pos += to.length();

}

/**

inline KernelCode(const std::string &pData = "", size_t workgroupSize = 256,

NumType precision = kf32)

: data(pData), workgroupSize({workgroupSize, 1, 1}),

precision(precision) {

if (precision == kf16) {

data = "enable f16;\n" + data;

}

replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1}));

replaceAll(data, "{{precision}}", toString(precision));

LOG(kDefLog, kTrace, "Shader code:\n%s", data.c_str());

}

/**

inline KernelCode(const std::string &pData,

const Shape &workgroupSize = {256, 1, 1},

NumType precision = kf32)

: data(pData), workgroupSize(workgroupSize), precision(precision) {

if (precision == kf16) {

data = "enable f16;\n" + data;

}

replaceAll(data, "{{workgroupSize}}", toString(workgroupSize));

replaceAll(data, "{{precision}}", toString(precision));

LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str());

}

/**

inline KernelCode(const std::string &pData, const Shape &workgroupSize,

NumType precision, const Shape &totalWorkgroups)

: data(pData), workgroupSize(workgroupSize), precision(precision) {

if (precision == kf16) {

data = "enable f16;\n" + data;

}

replaceAll(data, "{{workgroupSize}}", toString(workgroupSize));

replaceAll(data, "{{precision}}", toString(precision));

replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups));

LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str());

}

/**

inline KernelCode(const std::string &pData, const size_t &workgroupSize,

NumType precision, const Shape &totalWorkgroups)

: data(pData), workgroupSize({workgroupSize, 1, 1}),

precision(precision) {

if (precision == kf16) {

data = "enable f16;\n" + data;

}

replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1}));

replaceAll(data, "{{precision}}", toString(precision));

replaceAll(data, "{{totalWorkgroups}}", toString(totalWorkgroups));

LOG(kDefLog, kInfo, "Shader code:\n%s", data.c_str());

}

std::string data;

Shape workgroupSize;

NumType precision = kf32;

std::string label = "kernel";

std::string entryPoint = "main";

const std::vector<std::pair<std::string, std::string>> &reps) {

for (const auto &rep : reps) {

replaceAll(str, rep.first, rep.second);

}

return str;

WGPUBuffer buffer; // managed by owning Kernel

size_t bufferSize;

void *output; // non-owning, only for target memory in toCPU, not used for

// kernel invocations

std::promise<void> *promise;

std::future<void> *future;

WGPUCommandBuffer commandBuffer;

WGPUBuffer readbackBuffer;

std::promise<void> promise;

std::future<void> future;

std::unique_ptr<WGPUBuffer[]> buffers; // non-owning

std::unique_ptr<size_t[]> bufferSizes;

size_t numBindings;

Shape totalWorkgroups;

WGPUBindGroup bindGroup; // persists between submission

WGPUComputePipeline computePipeline; // persists between submission

WGPUCommandBuffer commandBuffer; // destroyed upon submission

bool used;

WGPUCompilationInfoRequestStatus status;

std::vector<std::string> messages;

std::vector<uint64_t> lineNums;

std::vector<uint64_t> linePos;

bool finished; // true if the compilation is finished

inline KernelPool(Context *ctx) : ctx(ctx), data() {}

Context *ctx;

std::unordered_map<std::string, Kernel> data;

inline ~KernelPool() {

// Note : Some kernel resources such as commandBuffer are harvested by

// queue submission, explicitly destroying readback and callback buffers

// produces runtime errors.

data.clear();

}

#ifdef __EMSCRIPTEN__

emscripten_sleep(0);

#else

wgpuInstanceProcessEvents(instance);

#endif

WGPUInstance instance = nullptr;

WGPUAdapter adapter = nullptr;

WGPUDevice device = nullptr;

WGPUQueue queue = nullptr;

TensorPool pool = TensorPool(this);

KernelPool kernelPool = KernelPool(this);

WGPURequestAdapterStatus adapterStatus;

WGPURequestDeviceStatus deviceStatus;

// Default constructor

Context() = default;

Context(Context&& other) noexcept

: instance(other.instance),

adapter(other.adapter),

device(other.device),

queue(other.queue),

// Re‐initialize pools to point to *this*:

pool(this),

kernelPool(this),

adapterStatus(other.adapterStatus),

deviceStatus(other.deviceStatus)

{

LOG(kDefLog, kTrace, "Moving Context ownership");

// Move over the resources in the pools:

pool.data = std::move(other.pool.data);

kernelPool.data = std::move(other.kernelPool.data);

// Null out handles in the source so its destructor won't release them.

other.instance = nullptr;

other.adapter = nullptr;

other.device = nullptr;

other.queue = nullptr;

// other.adapterStatus = 0;

// other.deviceStatus = 0;

}

Context& operator=(Context&& other) noexcept {

if (this != &other) {

// Free any existing resources. In most cases, this should be a no-op

// since we typically shouldn't have two active initialized Context

// instances with resources acquired.

this->~Context();

// Then placement‐new a move‐constructed copy in-place:

new (this) Context(std::move(other));

}

return *this;

}

~Context() {

LOG(kDefLog, kTrace, "Destroying context");

if (queue) {

wgpuQueueRelease(queue);

} else {

LOG(kDefLog, kTrace, "Queue is null");

}

if (device) {

wgpuDeviceRelease(device);

processEvents(instance);

} else {

LOG(kDefLog, kTrace, "Device is null");

}

if (adapter) {

wgpuAdapterRelease(adapter);

processEvents(instance);

} else {

LOG(kDefLog, kTrace, "Adapter is null");

}

if (instance) {

wgpuInstanceRelease(instance);

} else {

LOG(kDefLog, kTrace, "Instance is null");

}

LOG(kDefLog, kTrace, "Context destroyed");

}

const Shape &shape, NumType dtype,

WGPUBufferUsage usage = WGPUBufferUsage_Storage |

WGPUBufferUsage_CopyDst |

WGPUBufferUsage_CopySrc) {

LOG(kDefLog, kTrace, "Creating tensor");

size_t numElements = size(shape);

size_t size = sizeBytes(dtype) * numElements;

WGPUBufferDescriptor bufferDesc = {

.label = {.data = nullptr, .length = 0},

.usage = usage,

.size = size,

};

WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &bufferDesc);

pool.data[buffer] = Tensor{

.data = Array{.buffer = buffer, .usage = usage, .size = size},

.shape = shape,

};

return pool.data[buffer];

const float *data) {

assert(dtype == kf32);

Tensor tensor =

createTensor(ctx.pool, ctx.device, shape, dtype,

WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |

WGPUBufferUsage_CopySrc);

wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data,

tensor.data.size);

return tensor;

const int32_t *data) {

assert(dtype == ki32);

Tensor tensor =

createTensor(ctx.pool, ctx.device, shape, dtype,

WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |

WGPUBufferUsage_CopySrc);

wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data,

tensor.data.size);

return tensor;

const half *data) {

assert(dtype == kf16);

Tensor tensor =

createTensor(ctx.pool, ctx.device, shape, dtype,

WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |

WGPUBufferUsage_CopySrc);

wgpuQueueWriteBuffer(ctx.queue, tensor.data.buffer, 0, data,

tensor.data.size);

return tensor;

if (tensor.data.buffer) {

wgpuBufferRelease(tensor.data.buffer);

} else {

LOG(kDefLog, kWarn, "Tried to free tensor with null buffer");

}

if (pool.data.find(tensor.data.buffer) != pool.data.end()) {

pool.data.erase(tensor.data.buffer);

} else {

LOG(kDefLog, kWarn, "Tried to free tensor that was not in pool");

}

const char *file = "unkown", int line = -1) {

if (!condition) {

LOG(kDefLog, kError, "Error in file %s line %d:\n%s", file, line, message);

} else {

LOG(kDefLog, kTrace, "Success in file %s line %d:\n%s", file, line,

message);

}

const WGPUInstanceDescriptor &desc = {},

const WGPURequestAdapterOptions &adapterOpts = {},

const WGPUDeviceDescriptor &devDescriptor = {})

{

Context ctx; // stack-allocated

#ifdef __EMSCRIPTEN__

ctx.instance = wgpuCreateInstance(nullptr);

#else

ctx.instance = wgpuCreateInstance(&desc);

#endif

check(ctx.instance, "Initialize WebGPU", __FILE__, __LINE__);

LOG(kDefLog, kTrace, "Requesting adapter");

{

struct AdapterData {

WGPUAdapter adapter = nullptr;

bool requestEnded = false;

WGPURequestAdapterStatus status;

};

AdapterData adapterData;

auto onAdapterRequestEnded = [](WGPURequestAdapterStatus status,

WGPUAdapter adapter,

WGPUStringView message,

void *pUserData, void *) {

auto &ad = *reinterpret_cast<AdapterData*>(pUserData);

ad.status = status;

#ifdef __EMSCRIPTEN__

if (status != WGPURequestAdapterStatus_Success) {

LOG(kDefLog, kError, "Could not get WebGPU adapter: %.*s",

static_cast<int>(message.length), message.data);

}

#endif

check(status == WGPURequestAdapterStatus_Success,

"Request WebGPU adapter", __FILE__, __LINE__);

ad.adapter = adapter;

ad.requestEnded = true;

};

WGPURequestAdapterCallbackInfo callbackInfo {

.mode = WGPUCallbackMode_AllowSpontaneous,

.callback = onAdapterRequestEnded,

.userdata1 = &adapterData,

.userdata2 = nullptr

};

wgpuInstanceRequestAdapter(ctx.instance, &adapterOpts, callbackInfo);

while (!adapterData.requestEnded) {

processEvents(ctx.instance);

}

ctx.adapter = adapterData.adapter;

ctx.adapterStatus = adapterData.status;

}

LOG(kDefLog, kTrace, "Requesting device");

{

struct DeviceData {

WGPUDevice device = nullptr;

bool requestEnded = false;

WGPURequestDeviceStatus status;

};

DeviceData devData;

auto onDeviceRequestEnded = [](WGPURequestDeviceStatus status,

WGPUDevice device,

WGPUStringView message,

void *pUserData, void *) {

auto &dd = *reinterpret_cast<DeviceData*>(pUserData);

dd.status = status;

check(status == WGPURequestDeviceStatus_Success,

"Could not get WebGPU device.", __FILE__, __LINE__);

LOG(kDefLog, kTrace, "Device Request succeeded %p",

static_cast<void*>(device));

dd.device = device;

dd.requestEnded= true;

};

WGPURequestDeviceCallbackInfo deviceCallbackInfo {

.mode = WGPUCallbackMode_AllowSpontaneous,

.callback = onDeviceRequestEnded,

.userdata1= &devData,

.userdata2= nullptr

};

wgpuAdapterRequestDevice(ctx.adapter, &devDescriptor, deviceCallbackInfo);

LOG(kDefLog, kTrace, "Waiting for device request to end");

while (!devData.requestEnded) {

processEvents(ctx.instance);

}

LOG(kDefLog, kTrace, "Device request ended");

ctx.device = devData.device;

ctx.deviceStatus = devData.status;

// If the device was created, set up logging and fetch the queue

if (devData.status == WGPURequestDeviceStatus_Success) {

WGPULoggingCallbackInfo loggingCallbackInfo {

.nextInChain = nullptr,

.callback =

[](WGPULoggingType type, WGPUStringView message,

void *, void *) {

LOG(kDefLog, kError, "Device logging callback: %.*s",

static_cast<int>(message.length), message.data);

if (type == WGPULoggingType_Error) {

throw std::runtime_error("Device error logged.");

}

},

.userdata1 = nullptr,

.userdata2 = nullptr

};

wgpuDeviceSetLoggingCallback(ctx.device, loggingCallbackInfo);

ctx.queue = wgpuDeviceGetQueue(ctx.device);

}

}

return std::move(ctx);

const WGPUDeviceDescriptor &devDescriptor = {}) {

Context context;

{

#ifdef __EMSCRIPTEN__

// Emscripten does not support the instance descriptor

// and throws an assertion error if it is not nullptr.

context.instance = wgpuCreateInstance(nullptr);

#else

context.instance = wgpuCreateInstance(&desc);

#endif

// check status

check(context.instance, "Initialize WebGPU", __FILE__, __LINE__);

}

LOG(kDefLog, kInfo, "Requesting adapter");

{

std::vector<dawn::native::Adapter> adapters =

dawn::native::Instance(

reinterpret_cast<dawn::native::InstanceBase *>(context.instance))

.EnumerateAdapters();

LOG(kDefLog, kInfo, "The number of GPUs=%d\n", adapters.size());

// Note: Second gpu is not available on Macos, but the number of GPUs is 2

// on Macos.

// Calling wgpuAdapterGetInfo function for the second gpu becomes

// segfault. When you check all GPUs on linux, uncomment out following

// codes.

//

// for (size_t i = 0; i < adapters.size(); i++) {

// WGPUAdapterInfo info {};

// auto ptr = adapters[i].Get();

// if (ptr && adapters[i]) {

// wgpuAdapterGetInfo(ptr, &info);

// LOG(kDefLog, kInfo, "GPU(Adapter)[%d] = %s\n", i, info.description);

// wgpuAdapterInfoFreeMembers(info);

// }

// }

{

LOG(kDefLog, kInfo, "Use GPU(Adapter)[%d]\n", gpuIdx);

auto ptr = adapters[gpuIdx].Get();

if (ptr) {

WGPUAdapterInfo info{};