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/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