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%
6.6 KB · 159 lines cpp
Raw Blame History
1//-------------------------------------------------------------------------------------------------------------------------------------------------------------2//3// Metal/MTL4ComputePipeline.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 "MTL4PipelineState.hpp"25#include "MTLDefines.hpp"26#include "MTLHeaderBridge.hpp"27#include "MTLPrivate.hpp"28#include "MTLTypes.hpp"2930namespace MTL431{32class ComputePipelineDescriptor;33class FunctionDescriptor;34class StaticLinkingDescriptor;3536class ComputePipelineDescriptor : public NS::Copying<ComputePipelineDescriptor, PipelineDescriptor>37{38public:39    static ComputePipelineDescriptor* alloc();4041    FunctionDescriptor*               computeFunctionDescriptor() const;4243    ComputePipelineDescriptor*        init();4445    NS::UInteger                      maxTotalThreadsPerThreadgroup() const;4647    MTL::Size                         requiredThreadsPerThreadgroup() const;4849    void                              reset();5051    void                              setComputeFunctionDescriptor(const MTL4::FunctionDescriptor* computeFunctionDescriptor);5253    void                              setMaxTotalThreadsPerThreadgroup(NS::UInteger maxTotalThreadsPerThreadgroup);5455    void                              setRequiredThreadsPerThreadgroup(MTL::Size requiredThreadsPerThreadgroup);5657    void                              setStaticLinkingDescriptor(const MTL4::StaticLinkingDescriptor* staticLinkingDescriptor);5859    void                              setSupportBinaryLinking(bool supportBinaryLinking);6061    void                              setSupportIndirectCommandBuffers(MTL4::IndirectCommandBufferSupportState supportIndirectCommandBuffers);6263    void                              setThreadGroupSizeIsMultipleOfThreadExecutionWidth(bool threadGroupSizeIsMultipleOfThreadExecutionWidth);6465    StaticLinkingDescriptor*          staticLinkingDescriptor() const;6667    bool                              supportBinaryLinking() const;6869    IndirectCommandBufferSupportState supportIndirectCommandBuffers() const;7071    bool                              threadGroupSizeIsMultipleOfThreadExecutionWidth() const;72};7374}75_MTL_INLINE MTL4::ComputePipelineDescriptor* MTL4::ComputePipelineDescriptor::alloc()76{77    return NS::Object::alloc<MTL4::ComputePipelineDescriptor>(_MTL_PRIVATE_CLS(MTL4ComputePipelineDescriptor));78}7980_MTL_INLINE MTL4::FunctionDescriptor* MTL4::ComputePipelineDescriptor::computeFunctionDescriptor() const81{82    return Object::sendMessage<MTL4::FunctionDescriptor*>(this, _MTL_PRIVATE_SEL(computeFunctionDescriptor));83}8485_MTL_INLINE MTL4::ComputePipelineDescriptor* MTL4::ComputePipelineDescriptor::init()86{87    return NS::Object::init<MTL4::ComputePipelineDescriptor>();88}8990_MTL_INLINE NS::UInteger MTL4::ComputePipelineDescriptor::maxTotalThreadsPerThreadgroup() const91{92    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(maxTotalThreadsPerThreadgroup));93}9495_MTL_INLINE MTL::Size MTL4::ComputePipelineDescriptor::requiredThreadsPerThreadgroup() const96{97    return Object::sendMessage<MTL::Size>(this, _MTL_PRIVATE_SEL(requiredThreadsPerThreadgroup));98}99100_MTL_INLINE void MTL4::ComputePipelineDescriptor::reset()101{102    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(reset));103}104105_MTL_INLINE void MTL4::ComputePipelineDescriptor::setComputeFunctionDescriptor(const MTL4::FunctionDescriptor* computeFunctionDescriptor)106{107    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setComputeFunctionDescriptor_), computeFunctionDescriptor);108}109110_MTL_INLINE void MTL4::ComputePipelineDescriptor::setMaxTotalThreadsPerThreadgroup(NS::UInteger maxTotalThreadsPerThreadgroup)111{112    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setMaxTotalThreadsPerThreadgroup_), maxTotalThreadsPerThreadgroup);113}114115_MTL_INLINE void MTL4::ComputePipelineDescriptor::setRequiredThreadsPerThreadgroup(MTL::Size requiredThreadsPerThreadgroup)116{117    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setRequiredThreadsPerThreadgroup_), requiredThreadsPerThreadgroup);118}119120_MTL_INLINE void MTL4::ComputePipelineDescriptor::setStaticLinkingDescriptor(const MTL4::StaticLinkingDescriptor* staticLinkingDescriptor)121{122    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setStaticLinkingDescriptor_), staticLinkingDescriptor);123}124125_MTL_INLINE void MTL4::ComputePipelineDescriptor::setSupportBinaryLinking(bool supportBinaryLinking)126{127    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSupportBinaryLinking_), supportBinaryLinking);128}129130_MTL_INLINE void MTL4::ComputePipelineDescriptor::setSupportIndirectCommandBuffers(MTL4::IndirectCommandBufferSupportState supportIndirectCommandBuffers)131{132    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSupportIndirectCommandBuffers_), supportIndirectCommandBuffers);133}134135_MTL_INLINE void MTL4::ComputePipelineDescriptor::setThreadGroupSizeIsMultipleOfThreadExecutionWidth(bool threadGroupSizeIsMultipleOfThreadExecutionWidth)136{137    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setThreadGroupSizeIsMultipleOfThreadExecutionWidth_), threadGroupSizeIsMultipleOfThreadExecutionWidth);138}139140_MTL_INLINE MTL4::StaticLinkingDescriptor* MTL4::ComputePipelineDescriptor::staticLinkingDescriptor() const141{142    return Object::sendMessage<MTL4::StaticLinkingDescriptor*>(this, _MTL_PRIVATE_SEL(staticLinkingDescriptor));143}144145_MTL_INLINE bool MTL4::ComputePipelineDescriptor::supportBinaryLinking() const146{147    return Object::sendMessageSafe<bool>(this, _MTL_PRIVATE_SEL(supportBinaryLinking));148}149150_MTL_INLINE MTL4::IndirectCommandBufferSupportState MTL4::ComputePipelineDescriptor::supportIndirectCommandBuffers() const151{152    return Object::sendMessage<MTL4::IndirectCommandBufferSupportState>(this, _MTL_PRIVATE_SEL(supportIndirectCommandBuffers));153}154155_MTL_INLINE bool MTL4::ComputePipelineDescriptor::threadGroupSizeIsMultipleOfThreadExecutionWidth() const156{157    return Object::sendMessage<bool>(this, _MTL_PRIVATE_SEL(threadGroupSizeIsMultipleOfThreadExecutionWidth));158}159