MTLComputeTexture.hpp

Go to the documentation of this file.

#include <iostream>
#include <vector>
#include "MTLComputeGlobals.hpp"
 
#pragma once
 
namespace MTLCompute {
 
    template<typename T>
    class Texture {
        private:
            MTL::Device *gpu; 
            MTL::Texture *texture; 
            MTL::TextureDescriptor *descriptor; 
            int width = -1; 
            int height = -1; 
            bool freed = false; 
 
            void swap(Texture &tex) noexcept {
                using std::swap;
                swap(this->gpu, tex.gpu);
                swap(this->texture, tex.texture);
                swap(this->width, tex.width);
                swap(this->height, tex.height);
                swap(this->freed, tex.freed);
            }
 
            std::vector<T> flatten(std::vector<std::vector<T>> &v) const {
                std::vector<T> result;
                for (auto &inner : v) {
                    result.insert(result.end(), inner.begin(), inner.end());
                }
                return result;
            }
 
            std::vector<std::vector<T>> unflatten(std::vector<T> &v, int width, int height) const {
                std::vector<std::vector<T>> result(height, std::vector<T>(width));
                for (int i = 0; i < height; i++) {
                    for (int j = 0; j < width; j++) {
                        result[i][j] = v[i*width + j];
                    }
                }
                return result;
            }
            
        public:
 
            Texture(MTL::Device *gpu, int width, int height, TextureType tt) {
                this->gpu = gpu;
                if (width > MAX_TEXTURE_SIZE || height > MAX_TEXTURE_SIZE) {
                    throw std::invalid_argument("Texture size too large, max size is 16384");
                }
                this->width = width;
                this->height = height;
                this->descriptor = MTL::TextureDescriptor::alloc()->init();
                this->descriptor->setTextureType(MTL::TextureType2D);
                this->descriptor->setPixelFormat(static_cast<MTL::PixelFormat>(tt));
                if (sizeof(T) != TextureTypeSizes[tt]) {
                    std::cout << sizeof(T) << " " << TextureTypeSizes[tt] << std::endl;
                    throw std::invalid_argument("Texture type does not match template type");
                }
                this->descriptor->setWidth(width);
                this->descriptor->setHeight(height);
                this->texture = this->gpu->newTexture(this->descriptor);
            }
 
            Texture(MTL::Device *gpu, int width, int height) {
                this->gpu = gpu;
                if (width > MAX_TEXTURE_SIZE || height > MAX_TEXTURE_SIZE) {
                    throw std::invalid_argument("Texture size too large, max size is 16384");
                }
                this->width = width;
                this->height = height;
                this->descriptor = MTL::TextureDescriptor::alloc()->init();
                this->descriptor->setTextureType(MTL::TextureType2D);
                if (typeid(uint8_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR8Uint);
                } else if (typeid(uint16_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR16Uint);
                } else if (typeid(uint32_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR32Uint);
                } else if (typeid(float) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR32Float);
                } else if (typeid(int8_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR8Sint);
                } else if (typeid(int16_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR16Sint);
                } else if (typeid(int32_t) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR32Sint);
                } else if (typeid(float) == typeid(T)) {
                    this->descriptor->setPixelFormat(MTL::PixelFormatR32Float);
                } else {
                    throw std::invalid_argument("Texture type not supported");
                }
                this->descriptor->setWidth(width);
                this->descriptor->setHeight(height);
                this->texture = this->gpu->newTexture(this->descriptor);
            }
 
 
            Texture(const Texture &other) {
                this->gpu = other.gpu;
                this->texture = other.texture;
                this->width = other.width;
                this->height = other.height;
                this->descriptor = other.descriptor;
            }
 
 
            Texture() {
                this->gpu = nullptr;
                this->descriptor = nullptr;
                this->texture = nullptr;
                this->width = -1;
                this->height = -1;
            }
 
            ~Texture() {
                if (!this->freed) {
                    this->texture->autorelease();
                    this->descriptor->autorelease();
                    this->freed = true;
                }
            }
 
            void operator=(std::vector<std::vector<T>> data) {
                if (this->freed) {
                    throw std::runtime_error("Texture already freed");
                }
                if (data.size() != this->height || data[0].size() != this->width) {
                    throw std::invalid_argument("Data size does not match texture size");
                }
                if (data.size() > MAX_TEXTURE_SIZE || data[0].size() > MAX_TEXTURE_SIZE) {
                    throw std::invalid_argument("Data size too large, max size is 16384");
                }
                std::vector<T> flat = flatten(data);
                this->texture->replaceRegion(MTL::Region::Make2D(0, 0, this->width, this->height), 0, flat.data(), this->width*sizeof(T));
            }
 
            Texture & operator=(const Texture &other) {
 
                Texture temp(other);
                swap(temp);
 
                return *this;
            }
 
            std::vector<T> operator[](size_t index) const {
                if (this->freed) {
                    throw std::runtime_error("Texture already freed");
                }
                if (index >= (long)this->width*(long)this->height) {
                    throw std::out_of_range("Index out of bounds");
                }
                std::vector<T> flat((long)this->width*(long)this->height);
                this->texture->getBytes(flat.data(), this->width*sizeof(T), MTL::Region::Make2D(0, 0, this->width, this->height), 0);
                return unflatten(flat, this->width, this->height)[index];
            }
 
            std::vector<std::vector<T>> getData() {
                if (this->freed) {
                    throw std::runtime_error("Buffer already freed");
                }
                if (this->width == -1 || this->height == -1) {
                    throw std::runtime_error("Buffer not initialized");
                }
                std::vector<T> flat((long)this->width*(long)this->height);
                this->texture->getBytes(flat.data(), this->width*sizeof(T), MTL::Region::Make2D(0, 0, this->width, this->height), 0);
                return unflatten(flat, this->width, this->height);
            }
 
            MTL::Texture *getTexture() {
                return this->texture;
            }
 
            MTL::Device *getGPU() {
                return this->gpu;
            }
 
            bool getFreed() {
                return this->freed;
            }
 
            MTL::TextureDescriptor *getDescriptor() {
                return this->descriptor;
            }
 
            int getWidth() {
                return this->width;
            }
 
            int getHeight() {
                return this->height;
            }
    };
 
}