MTLComputeBuffer.hpp

Go to the documentation of this file.

#include <iostream>
#include <vector>
#include "MTLComputeGlobals.hpp"
 
#pragma once
 
namespace MTLCompute {
 
    template<typename T>
    class Buffer {
        private:
            MTL::Device *gpu; 
            MTL::Buffer *buffer; 
            bool freed = false; 
            MTLCompute::ResourceStorage storageMode; 
 
            void swap(Buffer &buffer) noexcept {
                using std::swap;
                swap(this->gpu, buffer.gpu);
                swap(this->buffer, buffer.buffer);
                swap(this->length, buffer.length);
                swap(this->itemsize, buffer.itemsize);
                swap(this->storageMode, buffer.storageMode);
            }
            
        public:
            Buffer(MTL::Device *gpu, size_t length, ResourceStorage storageMode) {
                this->gpu = gpu;
                this->length = length;
                this->itemsize = sizeof(T);
                this->storageMode = storageMode;
                this->buffer = gpu->newBuffer(length*itemsize, static_cast<MTL::ResourceOptions>(storageMode));
                this->buffer->retain();
            }
 
 
            Buffer(const Buffer &other) {
                this->gpu = other.gpu;
                this->length = other.length;
                this->itemsize = other.itemsize;
                this->storageMode = other.storageMode;
                this->buffer = other.buffer;
            }
 
 
            Buffer() {
                this->gpu = nullptr;
                this->length = -1;
                this->itemsize = -1;
                this->storageMode = MTLCompute::ResourceStorage::Shared;
                this->buffer = nullptr;
            }
 
            ~Buffer() {
                if (!this->freed) {
                    this->buffer->autorelease();
                    this->freed = true;
                }
            }
 
            T *contents() {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                return (T *)this->buffer->contents();
            }
 
            T operator[](size_t index) const {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                if (index >= this->length) {
                    throw std::out_of_range("Index out of bounds");
                }
                return ((T *)this->buffer->contents())[index];
            }
 
            T& operator[](size_t index) {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                if (index >= this->length) {
                    throw std::out_of_range("Index out of bounds");
                }
                return ((T *)this->buffer->contents())[index];
            }
 
            void operator=(std::vector<T> data) {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                if (data.size() != this->length) {
                    throw std::invalid_argument("Data size does not match buffer size");
                }
                memcpy(this->buffer->contents(), data.data(), this->length*this->itemsize);
 
                if (this->storageMode == MTLCompute::ResourceStorage::Managed) {
                    this->buffer->didModifyRange(NS::Range(0, this->length*this->itemsize));
                }
            }
 
            Buffer<T> & operator=(const Buffer &other) {
                // copy and swap because apparently it's good
                // i don't know what this is
                Buffer temp(other);
                swap(temp);
 
                return *this;
            }
 
            std::vector<T> getData() {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                if (this->length == -1) {
                    throw std::runtime_error("Buffer not initialized");
                }
                std::vector<T> data(this->length);
                memcpy(data.data(), this->buffer->contents(), this->length*this->itemsize);
                return data;
            }
 
            MTL::Buffer *getBuffer() {
                return this->buffer;
            }
 
            MTL::Device *getGPU() {
                return this->gpu;
            }
 
            bool getFreed() {
                return this->freed;
            }
 
            MTLCompute::ResourceStorage getStorageMode() {
                return this->storageMode;
            }
 
            size_t length; 
            size_t itemsize; 
    };
 
}