SPB Git

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%
16.8 KB · 440 lines cpp
Raw Blame History
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