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.4 KB · 173 lines cpp
Raw Blame History
1//-------------------------------------------------------------------------------------------------------------------------------------------------------------2//3// Metal/MTL4MachineLearningPipeline.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 "MTLAllocation.hpp"26#include "MTLDefines.hpp"27#include "MTLHeaderBridge.hpp"28#include "MTLPrivate.hpp"2930namespace MTL431{32class FunctionDescriptor;33class MachineLearningPipelineDescriptor;34class MachineLearningPipelineReflection;35}3637namespace MTL38{39class Device;40class TensorExtents;41}4243namespace MTL444{45class MachineLearningPipelineDescriptor : public NS::Copying<MachineLearningPipelineDescriptor, PipelineDescriptor>46{47public:48    static MachineLearningPipelineDescriptor* alloc();4950    MachineLearningPipelineDescriptor*        init();5152    MTL::TensorExtents*                       inputDimensionsAtBufferIndex(NS::Integer bufferIndex);5354    NS::String*                               label() const;5556    FunctionDescriptor*                       machineLearningFunctionDescriptor() const;5758    void                                      reset();5960    void                                      setInputDimensions(const MTL::TensorExtents* dimensions, NS::Integer bufferIndex);61    void                                      setInputDimensions(const NS::Array* dimensions, NS::Range range);6263    void                                      setLabel(const NS::String* label);6465    void                                      setMachineLearningFunctionDescriptor(const MTL4::FunctionDescriptor* machineLearningFunctionDescriptor);66};67class MachineLearningPipelineReflection : public NS::Referencing<MachineLearningPipelineReflection>68{69public:70    static MachineLearningPipelineReflection* alloc();7172    NS::Array*                                bindings() const;7374    MachineLearningPipelineReflection*        init();75};76class MachineLearningPipelineState : public NS::Referencing<MachineLearningPipelineState, MTL::Allocation>77{78public:79    MTL::Device*                       device() const;8081    NS::UInteger                       intermediatesHeapSize() const;8283    NS::String*                        label() const;8485    MachineLearningPipelineReflection* reflection() const;86};8788}89_MTL_INLINE MTL4::MachineLearningPipelineDescriptor* MTL4::MachineLearningPipelineDescriptor::alloc()90{91    return NS::Object::alloc<MTL4::MachineLearningPipelineDescriptor>(_MTL_PRIVATE_CLS(MTL4MachineLearningPipelineDescriptor));92}9394_MTL_INLINE MTL4::MachineLearningPipelineDescriptor* MTL4::MachineLearningPipelineDescriptor::init()95{96    return NS::Object::init<MTL4::MachineLearningPipelineDescriptor>();97}9899_MTL_INLINE MTL::TensorExtents* MTL4::MachineLearningPipelineDescriptor::inputDimensionsAtBufferIndex(NS::Integer bufferIndex)100{101    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(inputDimensionsAtBufferIndex_), bufferIndex);102}103104_MTL_INLINE NS::String* MTL4::MachineLearningPipelineDescriptor::label() const105{106    return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(label));107}108109_MTL_INLINE MTL4::FunctionDescriptor* MTL4::MachineLearningPipelineDescriptor::machineLearningFunctionDescriptor() const110{111    return Object::sendMessage<MTL4::FunctionDescriptor*>(this, _MTL_PRIVATE_SEL(machineLearningFunctionDescriptor));112}113114_MTL_INLINE void MTL4::MachineLearningPipelineDescriptor::reset()115{116    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(reset));117}118119_MTL_INLINE void MTL4::MachineLearningPipelineDescriptor::setInputDimensions(const MTL::TensorExtents* dimensions, NS::Integer bufferIndex)120{121    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setInputDimensions_atBufferIndex_), dimensions, bufferIndex);122}123124_MTL_INLINE void MTL4::MachineLearningPipelineDescriptor::setInputDimensions(const NS::Array* dimensions, NS::Range range)125{126    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setInputDimensions_withRange_), dimensions, range);127}128129_MTL_INLINE void MTL4::MachineLearningPipelineDescriptor::setLabel(const NS::String* label)130{131    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setLabel_), label);132}133134_MTL_INLINE void MTL4::MachineLearningPipelineDescriptor::setMachineLearningFunctionDescriptor(const MTL4::FunctionDescriptor* machineLearningFunctionDescriptor)135{136    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setMachineLearningFunctionDescriptor_), machineLearningFunctionDescriptor);137}138139_MTL_INLINE MTL4::MachineLearningPipelineReflection* MTL4::MachineLearningPipelineReflection::alloc()140{141    return NS::Object::alloc<MTL4::MachineLearningPipelineReflection>(_MTL_PRIVATE_CLS(MTL4MachineLearningPipelineReflection));142}143144_MTL_INLINE NS::Array* MTL4::MachineLearningPipelineReflection::bindings() const145{146    return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(bindings));147}148149_MTL_INLINE MTL4::MachineLearningPipelineReflection* MTL4::MachineLearningPipelineReflection::init()150{151    return NS::Object::init<MTL4::MachineLearningPipelineReflection>();152}153154_MTL_INLINE MTL::Device* MTL4::MachineLearningPipelineState::device() const155{156    return Object::sendMessage<MTL::Device*>(this, _MTL_PRIVATE_SEL(device));157}158159_MTL_INLINE NS::UInteger MTL4::MachineLearningPipelineState::intermediatesHeapSize() const160{161    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(intermediatesHeapSize));162}163164_MTL_INLINE NS::String* MTL4::MachineLearningPipelineState::label() const165{166    return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(label));167}168169_MTL_INLINE MTL4::MachineLearningPipelineReflection* MTL4::MachineLearningPipelineState::reflection() const170{171    return Object::sendMessage<MTL4::MachineLearningPipelineReflection*>(this, _MTL_PRIVATE_SEL(reflection));172}173