spb/forge Public MIT
Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.
C++ 61.2%
C 23%
Python 7.6%
TeX 7.2%
CMake 1.1%
1//-------------------------------------------------------------------------------------------------------------------------------------------------------------2//3// Metal/MTLComputePipeline.hpp4//5// Copyright 2020-2025 Apple Inc.6//7// Licensed under the Apache License, Version 2.0 (the "License");8// you may not use this file except in compliance with the License.9// You may obtain a copy of the License at10//11// http://www.apache.org/licenses/LICENSE-2.012//13// Unless required by applicable law or agreed to in writing, software14// distributed under the License is distributed on an "AS IS" BASIS,15// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.16// See the License for the specific language governing permissions and17// limitations under the License.18//19//-------------------------------------------------------------------------------------------------------------------------------------------------------------2021#pragma once2223#include "../Foundation/Foundation.hpp"24#include "MTLAllocation.hpp"25#include "MTLDefines.hpp"26#include "MTLHeaderBridge.hpp"27#include "MTLPipeline.hpp"28#include "MTLPrivate.hpp"29#include "MTLTypes.hpp"3031namespace MTL32{33class ComputePipelineDescriptor;34class ComputePipelineReflection;35class ComputePipelineState;36class Device;37class Function;38class FunctionHandle;39class IntersectionFunctionTable;40class IntersectionFunctionTableDescriptor;41class LinkedFunctions;42class PipelineBufferDescriptorArray;43class StageInputOutputDescriptor;44class VisibleFunctionTable;45class VisibleFunctionTableDescriptor;4647}48namespace MTL449{50class BinaryFunction;5152}53namespace MTL54{55class ComputePipelineReflection : public NS::Referencing<ComputePipelineReflection>56{57public:58 static ComputePipelineReflection* alloc();5960 NS::Array* arguments() const;6162 NS::Array* bindings() const;6364 ComputePipelineReflection* init();65};66class ComputePipelineDescriptor : public NS::Copying<ComputePipelineDescriptor>67{68public:69 static ComputePipelineDescriptor* alloc();7071 NS::Array* binaryArchives() const;7273 PipelineBufferDescriptorArray* buffers() const;7475 Function* computeFunction() const;7677 ComputePipelineDescriptor* init();7879 NS::Array* insertLibraries() const;8081 NS::String* label() const;8283 LinkedFunctions* linkedFunctions() const;8485 NS::UInteger maxCallStackDepth() const;8687 NS::UInteger maxTotalThreadsPerThreadgroup() const;8889 NS::Array* preloadedLibraries() const;9091 Size requiredThreadsPerThreadgroup() const;9293 void reset();9495 void setBinaryArchives(const NS::Array* binaryArchives);9697 void setComputeFunction(const MTL::Function* computeFunction);9899 void setInsertLibraries(const NS::Array* insertLibraries);100101 void setLabel(const NS::String* label);102103 void setLinkedFunctions(const MTL::LinkedFunctions* linkedFunctions);104105 void setMaxCallStackDepth(NS::UInteger maxCallStackDepth);106107 void setMaxTotalThreadsPerThreadgroup(NS::UInteger maxTotalThreadsPerThreadgroup);108109 void setPreloadedLibraries(const NS::Array* preloadedLibraries);110111 void setRequiredThreadsPerThreadgroup(MTL::Size requiredThreadsPerThreadgroup);112113 void setShaderValidation(MTL::ShaderValidation shaderValidation);114115 void setStageInputDescriptor(const MTL::StageInputOutputDescriptor* stageInputDescriptor);116117 void setSupportAddingBinaryFunctions(bool supportAddingBinaryFunctions);118119 void setSupportIndirectCommandBuffers(bool supportIndirectCommandBuffers);120121 void setThreadGroupSizeIsMultipleOfThreadExecutionWidth(bool threadGroupSizeIsMultipleOfThreadExecutionWidth);122123 ShaderValidation shaderValidation() const;124125 StageInputOutputDescriptor* stageInputDescriptor() const;126127 bool supportAddingBinaryFunctions() const;128129 bool supportIndirectCommandBuffers() const;130131 bool threadGroupSizeIsMultipleOfThreadExecutionWidth() const;132};133class ComputePipelineState : public NS::Referencing<ComputePipelineState, Allocation>134{135public:136 Device* device() const;137138 FunctionHandle* functionHandle(const NS::String* name);139 FunctionHandle* functionHandle(const MTL4::BinaryFunction* function);140 FunctionHandle* functionHandle(const MTL::Function* function);141142 ResourceID gpuResourceID() const;143144 NS::UInteger imageblockMemoryLength(MTL::Size imageblockDimensions);145146 NS::String* label() const;147148 NS::UInteger maxTotalThreadsPerThreadgroup() const;149150 ComputePipelineState* newComputePipelineStateWithBinaryFunctions(const NS::Array* additionalBinaryFunctions, NS::Error** error);151 ComputePipelineState* newComputePipelineState(const NS::Array* functions, NS::Error** error);152153 IntersectionFunctionTable* newIntersectionFunctionTable(const MTL::IntersectionFunctionTableDescriptor* descriptor);154155 VisibleFunctionTable* newVisibleFunctionTable(const MTL::VisibleFunctionTableDescriptor* descriptor);156157 ComputePipelineReflection* reflection() const;158159 Size requiredThreadsPerThreadgroup() const;160161 ShaderValidation shaderValidation() const;162163 NS::UInteger staticThreadgroupMemoryLength() const;164165 bool supportIndirectCommandBuffers() const;166167 NS::UInteger threadExecutionWidth() const;168};169170}171_MTL_INLINE MTL::ComputePipelineReflection* MTL::ComputePipelineReflection::alloc()172{173 return NS::Object::alloc<MTL::ComputePipelineReflection>(_MTL_PRIVATE_CLS(MTLComputePipelineReflection));174}175176_MTL_INLINE NS::Array* MTL::ComputePipelineReflection::arguments() const177{178 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(arguments));179}180181_MTL_INLINE NS::Array* MTL::ComputePipelineReflection::bindings() const182{183 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(bindings));184}185186_MTL_INLINE MTL::ComputePipelineReflection* MTL::ComputePipelineReflection::init()187{188 return NS::Object::init<MTL::ComputePipelineReflection>();189}190191_MTL_INLINE MTL::ComputePipelineDescriptor* MTL::ComputePipelineDescriptor::alloc()192{193 return NS::Object::alloc<MTL::ComputePipelineDescriptor>(_MTL_PRIVATE_CLS(MTLComputePipelineDescriptor));194}195196_MTL_INLINE NS::Array* MTL::ComputePipelineDescriptor::binaryArchives() const197{198 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(binaryArchives));199}200201_MTL_INLINE MTL::PipelineBufferDescriptorArray* MTL::ComputePipelineDescriptor::buffers() const202{203 return Object::sendMessage<MTL::PipelineBufferDescriptorArray*>(this, _MTL_PRIVATE_SEL(buffers));204}205206_MTL_INLINE MTL::Function* MTL::ComputePipelineDescriptor::computeFunction() const207{208 return Object::sendMessage<MTL::Function*>(this, _MTL_PRIVATE_SEL(computeFunction));209}210211_MTL_INLINE MTL::ComputePipelineDescriptor* MTL::ComputePipelineDescriptor::init()212{213 return NS::Object::init<MTL::ComputePipelineDescriptor>();214}215216_MTL_INLINE NS::Array* MTL::ComputePipelineDescriptor::insertLibraries() const217{218 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(insertLibraries));219}220221_MTL_INLINE NS::String* MTL::ComputePipelineDescriptor::label() const222{223 return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(label));224}225226_MTL_INLINE MTL::LinkedFunctions* MTL::ComputePipelineDescriptor::linkedFunctions() const227{228 return Object::sendMessage<MTL::LinkedFunctions*>(this, _MTL_PRIVATE_SEL(linkedFunctions));229}230231_MTL_INLINE NS::UInteger MTL::ComputePipelineDescriptor::maxCallStackDepth() const232{233 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(maxCallStackDepth));234}235236_MTL_INLINE NS::UInteger MTL::ComputePipelineDescriptor::maxTotalThreadsPerThreadgroup() const237{238 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(maxTotalThreadsPerThreadgroup));239}240241_MTL_INLINE NS::Array* MTL::ComputePipelineDescriptor::preloadedLibraries() const242{243 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(preloadedLibraries));244}245246_MTL_INLINE MTL::Size MTL::ComputePipelineDescriptor::requiredThreadsPerThreadgroup() const247{248 return Object::sendMessage<MTL::Size>(this, _MTL_PRIVATE_SEL(requiredThreadsPerThreadgroup));249}250251_MTL_INLINE void MTL::ComputePipelineDescriptor::reset()252{253 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(reset));254}255256_MTL_INLINE void MTL::ComputePipelineDescriptor::setBinaryArchives(const NS::Array* binaryArchives)257{258 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBinaryArchives_), binaryArchives);259}260261_MTL_INLINE void MTL::ComputePipelineDescriptor::setComputeFunction(const MTL::Function* computeFunction)262{263 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setComputeFunction_), computeFunction);264}265266_MTL_INLINE void MTL::ComputePipelineDescriptor::setInsertLibraries(const NS::Array* insertLibraries)267{268 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setInsertLibraries_), insertLibraries);269}270271_MTL_INLINE void MTL::ComputePipelineDescriptor::setLabel(const NS::String* label)272{273 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setLabel_), label);274}275276_MTL_INLINE void MTL::ComputePipelineDescriptor::setLinkedFunctions(const MTL::LinkedFunctions* linkedFunctions)277{278 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setLinkedFunctions_), linkedFunctions);279}280281_MTL_INLINE void MTL::ComputePipelineDescriptor::setMaxCallStackDepth(NS::UInteger maxCallStackDepth)282{283 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setMaxCallStackDepth_), maxCallStackDepth);284}285286_MTL_INLINE void MTL::ComputePipelineDescriptor::setMaxTotalThreadsPerThreadgroup(NS::UInteger maxTotalThreadsPerThreadgroup)287{288 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setMaxTotalThreadsPerThreadgroup_), maxTotalThreadsPerThreadgroup);289}290291_MTL_INLINE void MTL::ComputePipelineDescriptor::setPreloadedLibraries(const NS::Array* preloadedLibraries)292{293 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setPreloadedLibraries_), preloadedLibraries);294}295296_MTL_INLINE void MTL::ComputePipelineDescriptor::setRequiredThreadsPerThreadgroup(MTL::Size requiredThreadsPerThreadgroup)297{298 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setRequiredThreadsPerThreadgroup_), requiredThreadsPerThreadgroup);299}300301_MTL_INLINE void MTL::ComputePipelineDescriptor::setShaderValidation(MTL::ShaderValidation shaderValidation)302{303 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setShaderValidation_), shaderValidation);304}305306_MTL_INLINE void MTL::ComputePipelineDescriptor::setStageInputDescriptor(const MTL::StageInputOutputDescriptor* stageInputDescriptor)307{308 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setStageInputDescriptor_), stageInputDescriptor);309}310311_MTL_INLINE void MTL::ComputePipelineDescriptor::setSupportAddingBinaryFunctions(bool supportAddingBinaryFunctions)312{313 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSupportAddingBinaryFunctions_), supportAddingBinaryFunctions);314}315316_MTL_INLINE void MTL::ComputePipelineDescriptor::setSupportIndirectCommandBuffers(bool supportIndirectCommandBuffers)317{318 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSupportIndirectCommandBuffers_), supportIndirectCommandBuffers);319}320321_MTL_INLINE void MTL::ComputePipelineDescriptor::setThreadGroupSizeIsMultipleOfThreadExecutionWidth(bool threadGroupSizeIsMultipleOfThreadExecutionWidth)322{323 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setThreadGroupSizeIsMultipleOfThreadExecutionWidth_), threadGroupSizeIsMultipleOfThreadExecutionWidth);324}325326_MTL_INLINE MTL::ShaderValidation MTL::ComputePipelineDescriptor::shaderValidation() const327{328 return Object::sendMessage<MTL::ShaderValidation>(this, _MTL_PRIVATE_SEL(shaderValidation));329}330331_MTL_INLINE MTL::StageInputOutputDescriptor* MTL::ComputePipelineDescriptor::stageInputDescriptor() const332{333 return Object::sendMessage<MTL::StageInputOutputDescriptor*>(this, _MTL_PRIVATE_SEL(stageInputDescriptor));334}335336_MTL_INLINE bool MTL::ComputePipelineDescriptor::supportAddingBinaryFunctions() const337{338 return Object::sendMessageSafe<bool>(this, _MTL_PRIVATE_SEL(supportAddingBinaryFunctions));339}340341_MTL_INLINE bool MTL::ComputePipelineDescriptor::supportIndirectCommandBuffers() const342{343 return Object::sendMessageSafe<bool>(this, _MTL_PRIVATE_SEL(supportIndirectCommandBuffers));344}345346_MTL_INLINE bool MTL::ComputePipelineDescriptor::threadGroupSizeIsMultipleOfThreadExecutionWidth() const347{348 return Object::sendMessage<bool>(this, _MTL_PRIVATE_SEL(threadGroupSizeIsMultipleOfThreadExecutionWidth));349}350351_MTL_INLINE MTL::Device* MTL::ComputePipelineState::device() const352{353 return Object::sendMessage<MTL::Device*>(this, _MTL_PRIVATE_SEL(device));354}355356_MTL_INLINE MTL::FunctionHandle* MTL::ComputePipelineState::functionHandle(const NS::String* name)357{358 return Object::sendMessage<MTL::FunctionHandle*>(this, _MTL_PRIVATE_SEL(functionHandleWithName_), name);359}360361_MTL_INLINE MTL::FunctionHandle* MTL::ComputePipelineState::functionHandle(const MTL4::BinaryFunction* function)362{363 return Object::sendMessage<MTL::FunctionHandle*>(this, _MTL_PRIVATE_SEL(functionHandleWithBinaryFunction_), function);364}365366_MTL_INLINE MTL::FunctionHandle* MTL::ComputePipelineState::functionHandle(const MTL::Function* function)367{368 return Object::sendMessage<MTL::FunctionHandle*>(this, _MTL_PRIVATE_SEL(functionHandleWithFunction_), function);369}370371_MTL_INLINE MTL::ResourceID MTL::ComputePipelineState::gpuResourceID() const372{373 return Object::sendMessage<MTL::ResourceID>(this, _MTL_PRIVATE_SEL(gpuResourceID));374}375376_MTL_INLINE NS::UInteger MTL::ComputePipelineState::imageblockMemoryLength(MTL::Size imageblockDimensions)377{378 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(imageblockMemoryLengthForDimensions_), imageblockDimensions);379}380381_MTL_INLINE NS::String* MTL::ComputePipelineState::label() const382{383 return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(label));384}385386_MTL_INLINE NS::UInteger MTL::ComputePipelineState::maxTotalThreadsPerThreadgroup() const387{388 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(maxTotalThreadsPerThreadgroup));389}390391_MTL_INLINE MTL::ComputePipelineState* MTL::ComputePipelineState::newComputePipelineStateWithBinaryFunctions(const NS::Array* additionalBinaryFunctions, NS::Error** error)392{393 return Object::sendMessage<MTL::ComputePipelineState*>(this, _MTL_PRIVATE_SEL(newComputePipelineStateWithBinaryFunctions_error_), additionalBinaryFunctions, error);394}395396_MTL_INLINE MTL::ComputePipelineState* MTL::ComputePipelineState::newComputePipelineState(const NS::Array* functions, NS::Error** error)397{398 return Object::sendMessage<MTL::ComputePipelineState*>(this, _MTL_PRIVATE_SEL(newComputePipelineStateWithAdditionalBinaryFunctions_error_), functions, error);399}400401_MTL_INLINE MTL::IntersectionFunctionTable* MTL::ComputePipelineState::newIntersectionFunctionTable(const MTL::IntersectionFunctionTableDescriptor* descriptor)402{403 return Object::sendMessage<MTL::IntersectionFunctionTable*>(this, _MTL_PRIVATE_SEL(newIntersectionFunctionTableWithDescriptor_), descriptor);404}405406_MTL_INLINE MTL::VisibleFunctionTable* MTL::ComputePipelineState::newVisibleFunctionTable(const MTL::VisibleFunctionTableDescriptor* descriptor)407{408 return Object::sendMessage<MTL::VisibleFunctionTable*>(this, _MTL_PRIVATE_SEL(newVisibleFunctionTableWithDescriptor_), descriptor);409}410411_MTL_INLINE MTL::ComputePipelineReflection* MTL::ComputePipelineState::reflection() const412{413 return Object::sendMessage<MTL::ComputePipelineReflection*>(this, _MTL_PRIVATE_SEL(reflection));414}415416_MTL_INLINE MTL::Size MTL::ComputePipelineState::requiredThreadsPerThreadgroup() const417{418 return Object::sendMessage<MTL::Size>(this, _MTL_PRIVATE_SEL(requiredThreadsPerThreadgroup));419}420421_MTL_INLINE MTL::ShaderValidation MTL::ComputePipelineState::shaderValidation() const422{423 return Object::sendMessage<MTL::ShaderValidation>(this, _MTL_PRIVATE_SEL(shaderValidation));424}425426_MTL_INLINE NS::UInteger MTL::ComputePipelineState::staticThreadgroupMemoryLength() const427{428 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(staticThreadgroupMemoryLength));429}430431_MTL_INLINE bool MTL::ComputePipelineState::supportIndirectCommandBuffers() const432{433 return Object::sendMessageSafe<bool>(this, _MTL_PRIVATE_SEL(supportIndirectCommandBuffers));434}435436_MTL_INLINE NS::UInteger MTL::ComputePipelineState::threadExecutionWidth() const437{438 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(threadExecutionWidth));439}440