MTLComputeCommandManager.hpp

Go to the documentation of this file.

#include "MTLComputeBuffer.hpp"
#include "MTLComputeGlobals.hpp"
#include "MTLComputeKernel.hpp"
#include "MTLComputeTexture.hpp"
#include <iostream>
#include <vector>
 
#pragma once
 
namespace MTLCompute {
 
    template< typename T>
    class CommandManager {
 
        private:
            MTL::Device *gpu; 
            Kernel *kernel; 
            MTL::ComputePipelineState *pipeline; 
            MTL::CommandQueue *commandQueue; 
            MTL::CommandBuffer *commandBuffer; 
            MTL::ComputeCommandEncoder *commandEncoder; 
 
            std::vector<Buffer<T>> buffers = std::vector<Buffer<T>>(MAX_BUFFERS); 
            std::vector<Texture<T>> textures = std::vector<Texture<T>>(MAX_TEXTURES); 
            int bufferlength = -1; 
            int texwidth = -1; 
            int texheight = -1; 
 
        public:
 
            CommandManager(MTL::Device *gpu, MTLCompute::Kernel *kernel) {
                this->gpu = gpu;
                this->kernel = kernel;
                this->pipeline = this->kernel->getPLS();
 
                this->commandQueue = this->gpu->newCommandQueue();
 
            }
 
            CommandManager() = default;
 
            ~CommandManager() {
                this->commandQueue->autorelease();
            }
 
            void loadBuffer(Buffer<T> buffer, int index) {
                if (this->bufferlength == -1) {
                    this->bufferlength = buffer.length;
                } else if (this->bufferlength != buffer.length) {
                    throw std::invalid_argument("Buffer lengths do not match");
                }
 
                this->buffers[index] = buffer;
            }
 
 
            void loadTexture(Texture<T> texture, int index) {
                if (this->texwidth == -1 && this->texheight == -1) {
                    this->texwidth = texture.getWidth();
                    this->texheight = texture.getHeight();
                    if (this->texwidth > MAX_TEXTURE_SIZE || this->texheight > MAX_TEXTURE_SIZE) {
                        throw std::invalid_argument("Texture size too large, max size is 16384");
                    }
                } else if (this->texwidth != texture.getWidth() || this->texheight != texture.getHeight()) {
                    std::cout << this->texwidth << " " << this->texheight << std::endl;
                    std::cout << texture.getTexture()->width() << " " << texture.getTexture()->height() << std::endl;
                    throw std::invalid_argument("Texture sizes do not match");
                }
 
                this->textures[index] = texture;
            }
 
            void dispatch() {
                if (this->kernel->getPLS() != this->pipeline) {
                    // Refresh the pipeline if it has changed
                    this->pipeline = this->kernel->getPLS();
                }
 
                // Create a new command buffer and command encoder
                this->commandBuffer = this->commandQueue->commandBuffer();
                this->commandEncoder = this->commandBuffer->computeCommandEncoder();
                this->commandEncoder->setComputePipelineState(this->pipeline);
                bool usingbuffers = false;
                bool usingtextures = false;
 
                // Load the buffers and textures into the commandEncoder
                for (int i = 0; i < MAX_BUFFERS; i++) {
                    if (buffers[i].length == this->bufferlength && buffers[i].getBuffer() != nullptr) {
                        this->commandEncoder->setBuffer(buffers[i].getBuffer(), 0, i);
                        usingbuffers = true;
                    }
                    if (textures[i].getWidth() == this->texwidth && textures[i].getHeight() == this->texheight
                            && textures[i].getTexture() != nullptr) {
                        this->commandEncoder->setTexture(textures[i].getTexture(), i);
                        usingtextures = true;
                    }
                }
                for (int i = MAX_BUFFERS; i < MAX_TEXTURES; i++) {
                    if (textures[i].getWidth() == this->texwidth && textures[i].getHeight() == this->texheight
                            && textures[i].getTexture() != nullptr) {
                        this->commandEncoder->setTexture(textures[i].getTexture(), i);
                        usingtextures = true;
                    }
                }
 
                // Calculate the grid size and thread group size
                MTL::Size threadsPerThreadgroup;
                threadsPerThreadgroup.width = this->pipeline->threadExecutionWidth();
                threadsPerThreadgroup.height = this->pipeline->maxTotalThreadsPerThreadgroup() / threadsPerThreadgroup.width;
                threadsPerThreadgroup.depth = 1;
 
                MTL::Size threadsPerGrid;
                if (usingbuffers && usingtextures) {
                    if (this->bufferlength > this->texwidth)
                        threadsPerGrid = MTL::Size::Make(this->bufferlength, this->texheight, 1);
                    else
                        threadsPerGrid = MTL::Size::Make(this->texwidth, this->texheight, 1);
                } else if (usingbuffers && !usingtextures) {
                    threadsPerGrid = MTL::Size::Make(this->bufferlength, 1, 1);
 
                } else if (!usingbuffers && usingtextures) {
                    threadsPerGrid = MTL::Size::Make(this->texwidth, this->texheight, 1);
                    
                } else {
                    throw std::invalid_argument("No buffers or textures loaded");
                }
 
                // Use dispatchThreads NOT dispatchThreadgroups
                this->commandEncoder->dispatchThreads(threadsPerGrid, threadsPerThreadgroup);
                this->commandEncoder->endEncoding();
                this->commandBuffer->commit();
                this->commandBuffer->waitUntilCompleted();
 
                // Release the command encoder and command buffer
                this->commandEncoder->release();
                this->commandBuffer->release();
            }
 
            void resetBuffers() {
                this->buffers.clear();
                this->buffers = std::vector<Buffer<T>>(MAX_BUFFERS);
                this->bufferlength = -1;
            }
 
            void resetTextures() {
                this->textures.clear();
                this->textures = std::vector<Texture<T>>(MAX_TEXTURES);
                this->texwidth = -1;
                this->texheight = -1;
            }
 
            void reset() {
                this->resetBuffers();
                this->resetTextures();
            }
 
            MTL::Device *getGPU() {
                return this->gpu;
            }
 
            Kernel *getKernel() {
                return this->kernel;
            }
 
            std::vector<Buffer<T>>& getBuffers() {
                return this->buffers;
            }
 
            std::vector<Texture<T>>& getTextures() {
                return this->textures;
            }
 
    };
 
}