WGPUBuffer buffer;

WGPUBufferUsageFlags 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::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();

switch (type) {

case kf16:

return sizeof(uint16_t);

case kf32:

return sizeof(float);

default:

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

return 0;

}

switch (type) {

case kf16:

return "f16";

case kf32:

return "f32";

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) {

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

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

LOG(kDefLog, kInfo, "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) {

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

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

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);

}

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;

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

std::unique_ptr<size_t[]> bufferSizes;

size_t numBindings;

Shape nWorkgroups;

WGPUBindGroup bindGroup; // persists between submission

WGPUComputePipeline computePipeline; // persists between submission

WGPUCommandBuffer commandBuffer; // destroyed upon submission

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

Context *ctx;

std::set<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();

}

WGPUInstance instance;

WGPUAdapter adapter;

WGPUDevice device;

WGPUQueue queue;

TensorPool pool = TensorPool(this);

KernelPool kernelPool = KernelPool(this);

~Context() {

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

if (queue) {

wgpuQueueRelease(queue);

wgpuInstanceProcessEvents(instance);

} else {

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

}

if (device) {

wgpuDeviceRelease(device);

wgpuInstanceProcessEvents(instance);

} else {

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

}

if (adapter) {

wgpuAdapterRelease(adapter);

wgpuInstanceProcessEvents(instance);

} else {

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

}

if (instance) {

wgpuInstanceRelease(instance);

} else {

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

}

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

}

NumType dtype,

WGPUBufferUsageFlags usage = WGPUBufferUsage_Storage |

WGPUBufferUsage_CopyDst |

WGPUBufferUsage_CopySrc) {

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

size_t numElements = 1;

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

numElements *= shape.data[dim];

}

size_t size = dtype == kf32 ? sizeBytes(dtype) * numElements : 0;

WGPUBufferDescriptor bufferDesc = {

.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];

float *data) {

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 constexpr (kDebug) {

if (!condition) {

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

message);

exit(1);

} else {

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

message);

}

}

const WGPURequestAdapterOptions &adapterOpts = {},

WGPUDeviceDescriptor devDescriptor = {}) {

Context context;

{

context.instance = wgpuCreateInstance(&desc);

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

}

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

{

struct AdapterData {

WGPUAdapter adapter = nullptr;

bool requestEnded = false;

};

AdapterData adapterData;

auto onAdapterRequestEnded = [](WGPURequestAdapterStatus status,

WGPUAdapter adapter, char const *message,

void *pUserData) {

AdapterData &adapterData = *reinterpret_cast<AdapterData *>(pUserData);

check(status == WGPURequestAdapterStatus_Success,

"Request WebGPU adapter", __FILE__, __LINE__);

adapterData.adapter = adapter;

adapterData.requestEnded = true;

};

wgpuInstanceRequestAdapter(context.instance, &adapterOpts,

onAdapterRequestEnded, (void *)&adapterData);

assert(adapterData.requestEnded);

context.adapter = adapterData.adapter;

}

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

{

struct DeviceData {

WGPUDevice device = nullptr;

bool requestEnded = false;

};

DeviceData devData;

auto onDeviceRequestEnded = [](WGPURequestDeviceStatus status,

WGPUDevice device, char const *message,

void *pUserData) {

DeviceData &devData = *reinterpret_cast<DeviceData *>(pUserData);

check(status == WGPURequestDeviceStatus_Success,

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

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

static_cast<void *>(device));

devData.device = device;

devData.requestEnded = true;

};

#ifdef WEBGPU_BACKEND_DAWN

devDescriptor.deviceLostCallbackInfo = {

.callback =

[](WGPUDevice const *device, WGPUDeviceLostReason reason,

char const *message, void *userdata) {

if (reason != WGPUDeviceLostReason_Destroyed) {

LOG(kDefLog, kError, "Device lost (code %d):\n%s", reason,

message);

} else {

LOG(kDefLog, kInfo, "Device destroyed: %s", message);

}

},

};

#endif

wgpuAdapterRequestDevice(context.adapter, &devDescriptor,

onDeviceRequestEnded, (void *)&devData);

assert(devData.requestEnded);

context.device = devData.device;

wgpuDeviceSetUncapturedErrorCallback(

context.device,

[](WGPUErrorType type, char const *message, void *devData) {

LOG(kDefLog, kError, "Device uncaptured error: %s", message);

throw std::runtime_error("Device uncaptured exception.");

},

nullptr);

}

context.queue = wgpuDeviceGetQueue(context.device);

return context;

while (future.wait_for(std::chrono::seconds(0)) !=

std::future_status::ready) {

wgpuInstanceProcessEvents(ctx.instance);

}

WGPUDevice device = ctx.device;

struct CopyOp {

WGPUCommandBuffer commandBuffer;

WGPUBuffer readbackBuffer;

std::promise<void> promise;

std::future<void> future;

CallbackData callbackData;

};

CopyOp op;

op.future = op.promise.get_future();

{

WGPUBufferDescriptor readbackBufferDescriptor = {

.usage = WGPUBufferUsage_CopyDst | WGPUBufferUsage_MapRead,

.size = bufferSize,

};

op.readbackBuffer =

wgpuDeviceCreateBuffer(device, &readbackBufferDescriptor);

}

{

WGPUCommandEncoder commandEncoder;

WGPUComputePassEncoder computePassEncoder;

commandEncoder = wgpuDeviceCreateCommandEncoder(device, nullptr);

wgpuCommandEncoderCopyBufferToBuffer(commandEncoder, tensor.data.buffer, 0,

op.readbackBuffer, 0, bufferSize);

op.commandBuffer = wgpuCommandEncoderFinish(commandEncoder, nullptr);

check(op.commandBuffer, "Create command buffer", __FILE__, __LINE__);

}

wgpuQueueSubmit(ctx.queue, 1, &op.commandBuffer);

CallbackData callbackData = {op.readbackBuffer, bufferSize, data, &op.promise,

&op.future};

wgpuQueueOnSubmittedWorkDone(

ctx.queue,

[](WGPUQueueWorkDoneStatus status, void *callbackData) {

check(status == WGPUQueueWorkDoneStatus_Success, "Queue work done",

__FILE__, __LINE__);

const auto *data = static_cast<CallbackData *>(callbackData);

wgpuBufferMapAsync(

data->buffer, WGPUMapMode_Read, 0, data->bufferSize,

[](WGPUBufferMapAsyncStatus status, void *captureData) {

const auto *data = static_cast<CallbackData *>(captureData);

check(status == WGPUBufferMapAsyncStatus_Success,

"Map readbackBuffer", __FILE__, __LINE__);

const void *mappedData = wgpuBufferGetConstMappedRange(

data->buffer, /*offset=*/0, data->bufferSize);

check(mappedData, "Get mapped range", __FILE__, __LINE__);

memcpy(data->output, mappedData, data->bufferSize);

wgpuBufferUnmap(data->buffer);

data->promise->set_value();

},

callbackData);

},

&callbackData);

wait(ctx, op.future);

// TODO(avh): Maintain params metadata in Kernel and check for consistency.

// If a kernel does not have parameters this will quietly overwrite

// the last buffer in the bind group with the parameters buffer.

if (op.numBindings > 0) {

wgpuQueueWriteBuffer(ctx.queue, op.buffers[op.numBindings - 1], 0,

static_cast<void *>(&params), sizeof(params));

}

{

WGPUCommandEncoder commandEncoder =

wgpuDeviceCreateCommandEncoder(device, nullptr);

WGPUComputePassEncoder computePassEncoder =

wgpuCommandEncoderBeginComputePass(commandEncoder, nullptr);

wgpuComputePassEncoderSetPipeline(computePassEncoder, op.computePipeline);

wgpuComputePassEncoderSetBindGroup(computePassEncoder, 0, op.bindGroup, 0,

nullptr);

wgpuComputePassEncoderDispatchWorkgroups(

computePassEncoder, op.nWorkgroups[0], op.nWorkgroups[1],

op.nWorkgroups[2]);

wgpuComputePassEncoderEnd(computePassEncoder);

op.commandBuffer = wgpuCommandEncoderFinish(commandEncoder, nullptr);

}

assert(total.rank == group.rank);

Shape result;

result.rank = total.rank;

for (size_t dim = 0; dim < total.rank; ++dim) {

result[dim] = cdiv(total[dim], group[dim]);

}

return result;

const Tensor *dataBindings, size_t numTensors,

const size_t *viewOffsets, const Shape &nWorkgroups,

const void *params = nullptr,

size_t paramsSize = 0) {

assert(nWorkgroups.rank == 3);

WGPUDevice device = ctx.device;

WGPUQueue queue = ctx.queue;

Kernel op;

// paramIndex is the index into bgLayoutEntries for the parameters buffer If

// there are no parameters for the kernel, paramsSize == 0 and paramIndex is

// effectively undefined (== -1)

size_t paramIndex = -1;

// Note: paramIndex is undefined unless paramsSize > 0

size_t numBindings = numTensors;

if (paramsSize > 0) {

numBindings++; // parameters buffer

paramIndex = numBindings - 1; // index of the parameters buffer within

// op.buffers, op.bufferSizes and

// bgLayoutEntries

}

op.buffers = std::make_unique<WGPUBuffer[]>(numBindings);

op.bufferSizes = std::make_unique<size_t[]>(numBindings);

op.numBindings = numBindings;

std::vector<WGPUBindGroupLayoutEntry> bgLayoutEntries(numBindings);

// Create layout entries for input buffers

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

bgLayoutEntries[i] = WGPUBindGroupLayoutEntry{

.binding = static_cast<uint32_t>(i),

.visibility = WGPUShaderStage_Compute,

.buffer =

WGPUBufferBindingLayout{

.type = WGPUBufferBindingType_Storage,

.minBindingSize = dataBindings[i].data.size,

},

};

}

if (paramsSize > 0) {

LOG(kDefLog, kInfo, "Create layout entry for the params buffer");

// Create layout entry for the params buffer

bgLayoutEntries[paramIndex] = WGPUBindGroupLayoutEntry{

.binding = static_cast<uint32_t>(paramIndex),

.visibility = WGPUShaderStage_Compute,

.buffer =

WGPUBufferBindingLayout{

.type = WGPUBufferBindingType_Uniform,

.minBindingSize = paramsSize,

},

};

}

WGPUBindGroupLayoutDescriptor bgLayoutDesc = {

.entryCount = static_cast<uint32_t>(bgLayoutEntries.size()),

.entries = bgLayoutEntries.data(),

};

WGPUBindGroupLayout bgLayout =

wgpuDeviceCreateBindGroupLayout(device, &bgLayoutDesc);

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

op.buffers[i] = dataBindings[i].data.buffer;

op.bufferSizes[i] = dataBindings[i].data.size;

}

// Create a buffer for the Params struct

if (paramsSize > 0) {

WGPUBufferDescriptor paramsBufferDesc = {

.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst,

.size = paramsSize,

.mappedAtCreation = false,