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/MTLFunctionStitching.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 "MTLDefines.hpp"25#include "MTLHeaderBridge.hpp"26#include "MTLPrivate.hpp"2728namespace MTL29{30class FunctionStitchingAttributeAlwaysInline;31class FunctionStitchingFunctionNode;32class FunctionStitchingGraph;33class FunctionStitchingInputNode;34class StitchedLibraryDescriptor;3536_MTL_OPTIONS(NS::UInteger, StitchedLibraryOptions) {37 StitchedLibraryOptionNone = 0,38 StitchedLibraryOptionFailOnBinaryArchiveMiss = 1,39 StitchedLibraryOptionStoreLibraryInMetalPipelinesScript = 1 << 1,40};4142class FunctionStitchingAttribute : public NS::Referencing<FunctionStitchingAttribute>43{44};45class FunctionStitchingAttributeAlwaysInline : public NS::Referencing<FunctionStitchingAttributeAlwaysInline, FunctionStitchingAttribute>46{47public:48 static FunctionStitchingAttributeAlwaysInline* alloc();4950 FunctionStitchingAttributeAlwaysInline* init();51};52class FunctionStitchingNode : public NS::Copying<FunctionStitchingNode>53{54};55class FunctionStitchingInputNode : public NS::Referencing<FunctionStitchingInputNode, FunctionStitchingNode>56{57public:58 static FunctionStitchingInputNode* alloc();5960 NS::UInteger argumentIndex() const;6162 FunctionStitchingInputNode* init();63 FunctionStitchingInputNode* init(NS::UInteger argument);6465 void setArgumentIndex(NS::UInteger argumentIndex);66};67class FunctionStitchingFunctionNode : public NS::Referencing<FunctionStitchingFunctionNode, FunctionStitchingNode>68{69public:70 static FunctionStitchingFunctionNode* alloc();7172 NS::Array* arguments() const;7374 NS::Array* controlDependencies() const;7576 FunctionStitchingFunctionNode* init();77 FunctionStitchingFunctionNode* init(const NS::String* name, const NS::Array* arguments, const NS::Array* controlDependencies);7879 NS::String* name() const;8081 void setArguments(const NS::Array* arguments);8283 void setControlDependencies(const NS::Array* controlDependencies);8485 void setName(const NS::String* name);86};87class FunctionStitchingGraph : public NS::Copying<FunctionStitchingGraph>88{89public:90 static FunctionStitchingGraph* alloc();9192 NS::Array* attributes() const;9394 NS::String* functionName() const;9596 FunctionStitchingGraph* init();97 FunctionStitchingGraph* init(const NS::String* functionName, const NS::Array* nodes, const MTL::FunctionStitchingFunctionNode* outputNode, const NS::Array* attributes);9899 NS::Array* nodes() const;100101 FunctionStitchingFunctionNode* outputNode() const;102103 void setAttributes(const NS::Array* attributes);104105 void setFunctionName(const NS::String* functionName);106107 void setNodes(const NS::Array* nodes);108109 void setOutputNode(const MTL::FunctionStitchingFunctionNode* outputNode);110};111class StitchedLibraryDescriptor : public NS::Copying<StitchedLibraryDescriptor>112{113public:114 static StitchedLibraryDescriptor* alloc();115116 NS::Array* binaryArchives() const;117118 NS::Array* functionGraphs() const;119120 NS::Array* functions() const;121122 StitchedLibraryDescriptor* init();123124 StitchedLibraryOptions options() const;125126 void setBinaryArchives(const NS::Array* binaryArchives);127128 void setFunctionGraphs(const NS::Array* functionGraphs);129130 void setFunctions(const NS::Array* functions);131132 void setOptions(MTL::StitchedLibraryOptions options);133};134135}136_MTL_INLINE MTL::FunctionStitchingAttributeAlwaysInline* MTL::FunctionStitchingAttributeAlwaysInline::alloc()137{138 return NS::Object::alloc<MTL::FunctionStitchingAttributeAlwaysInline>(_MTL_PRIVATE_CLS(MTLFunctionStitchingAttributeAlwaysInline));139}140141_MTL_INLINE MTL::FunctionStitchingAttributeAlwaysInline* MTL::FunctionStitchingAttributeAlwaysInline::init()142{143 return NS::Object::init<MTL::FunctionStitchingAttributeAlwaysInline>();144}145146_MTL_INLINE MTL::FunctionStitchingInputNode* MTL::FunctionStitchingInputNode::alloc()147{148 return NS::Object::alloc<MTL::FunctionStitchingInputNode>(_MTL_PRIVATE_CLS(MTLFunctionStitchingInputNode));149}150151_MTL_INLINE NS::UInteger MTL::FunctionStitchingInputNode::argumentIndex() const152{153 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(argumentIndex));154}155156_MTL_INLINE MTL::FunctionStitchingInputNode* MTL::FunctionStitchingInputNode::init()157{158 return NS::Object::init<MTL::FunctionStitchingInputNode>();159}160161_MTL_INLINE MTL::FunctionStitchingInputNode* MTL::FunctionStitchingInputNode::init(NS::UInteger argument)162{163 return Object::sendMessage<MTL::FunctionStitchingInputNode*>(this, _MTL_PRIVATE_SEL(initWithArgumentIndex_), argument);164}165166_MTL_INLINE void MTL::FunctionStitchingInputNode::setArgumentIndex(NS::UInteger argumentIndex)167{168 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setArgumentIndex_), argumentIndex);169}170171_MTL_INLINE MTL::FunctionStitchingFunctionNode* MTL::FunctionStitchingFunctionNode::alloc()172{173 return NS::Object::alloc<MTL::FunctionStitchingFunctionNode>(_MTL_PRIVATE_CLS(MTLFunctionStitchingFunctionNode));174}175176_MTL_INLINE NS::Array* MTL::FunctionStitchingFunctionNode::arguments() const177{178 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(arguments));179}180181_MTL_INLINE NS::Array* MTL::FunctionStitchingFunctionNode::controlDependencies() const182{183 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(controlDependencies));184}185186_MTL_INLINE MTL::FunctionStitchingFunctionNode* MTL::FunctionStitchingFunctionNode::init()187{188 return NS::Object::init<MTL::FunctionStitchingFunctionNode>();189}190191_MTL_INLINE MTL::FunctionStitchingFunctionNode* MTL::FunctionStitchingFunctionNode::init(const NS::String* name, const NS::Array* arguments, const NS::Array* controlDependencies)192{193 return Object::sendMessage<MTL::FunctionStitchingFunctionNode*>(this, _MTL_PRIVATE_SEL(initWithName_arguments_controlDependencies_), name, arguments, controlDependencies);194}195196_MTL_INLINE NS::String* MTL::FunctionStitchingFunctionNode::name() const197{198 return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(name));199}200201_MTL_INLINE void MTL::FunctionStitchingFunctionNode::setArguments(const NS::Array* arguments)202{203 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setArguments_), arguments);204}205206_MTL_INLINE void MTL::FunctionStitchingFunctionNode::setControlDependencies(const NS::Array* controlDependencies)207{208 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setControlDependencies_), controlDependencies);209}210211_MTL_INLINE void MTL::FunctionStitchingFunctionNode::setName(const NS::String* name)212{213 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setName_), name);214}215216_MTL_INLINE MTL::FunctionStitchingGraph* MTL::FunctionStitchingGraph::alloc()217{218 return NS::Object::alloc<MTL::FunctionStitchingGraph>(_MTL_PRIVATE_CLS(MTLFunctionStitchingGraph));219}220221_MTL_INLINE NS::Array* MTL::FunctionStitchingGraph::attributes() const222{223 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(attributes));224}225226_MTL_INLINE NS::String* MTL::FunctionStitchingGraph::functionName() const227{228 return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(functionName));229}230231_MTL_INLINE MTL::FunctionStitchingGraph* MTL::FunctionStitchingGraph::init()232{233 return NS::Object::init<MTL::FunctionStitchingGraph>();234}235236_MTL_INLINE MTL::FunctionStitchingGraph* MTL::FunctionStitchingGraph::init(const NS::String* functionName, const NS::Array* nodes, const MTL::FunctionStitchingFunctionNode* outputNode, const NS::Array* attributes)237{238 return Object::sendMessage<MTL::FunctionStitchingGraph*>(this, _MTL_PRIVATE_SEL(initWithFunctionName_nodes_outputNode_attributes_), functionName, nodes, outputNode, attributes);239}240241_MTL_INLINE NS::Array* MTL::FunctionStitchingGraph::nodes() const242{243 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(nodes));244}245246_MTL_INLINE MTL::FunctionStitchingFunctionNode* MTL::FunctionStitchingGraph::outputNode() const247{248 return Object::sendMessage<MTL::FunctionStitchingFunctionNode*>(this, _MTL_PRIVATE_SEL(outputNode));249}250251_MTL_INLINE void MTL::FunctionStitchingGraph::setAttributes(const NS::Array* attributes)252{253 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setAttributes_), attributes);254}255256_MTL_INLINE void MTL::FunctionStitchingGraph::setFunctionName(const NS::String* functionName)257{258 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setFunctionName_), functionName);259}260261_MTL_INLINE void MTL::FunctionStitchingGraph::setNodes(const NS::Array* nodes)262{263 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setNodes_), nodes);264}265266_MTL_INLINE void MTL::FunctionStitchingGraph::setOutputNode(const MTL::FunctionStitchingFunctionNode* outputNode)267{268 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setOutputNode_), outputNode);269}270271_MTL_INLINE MTL::StitchedLibraryDescriptor* MTL::StitchedLibraryDescriptor::alloc()272{273 return NS::Object::alloc<MTL::StitchedLibraryDescriptor>(_MTL_PRIVATE_CLS(MTLStitchedLibraryDescriptor));274}275276_MTL_INLINE NS::Array* MTL::StitchedLibraryDescriptor::binaryArchives() const277{278 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(binaryArchives));279}280281_MTL_INLINE NS::Array* MTL::StitchedLibraryDescriptor::functionGraphs() const282{283 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(functionGraphs));284}285286_MTL_INLINE NS::Array* MTL::StitchedLibraryDescriptor::functions() const287{288 return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(functions));289}290291_MTL_INLINE MTL::StitchedLibraryDescriptor* MTL::StitchedLibraryDescriptor::init()292{293 return NS::Object::init<MTL::StitchedLibraryDescriptor>();294}295296_MTL_INLINE MTL::StitchedLibraryOptions MTL::StitchedLibraryDescriptor::options() const297{298 return Object::sendMessage<MTL::StitchedLibraryOptions>(this, _MTL_PRIVATE_SEL(options));299}300301_MTL_INLINE void MTL::StitchedLibraryDescriptor::setBinaryArchives(const NS::Array* binaryArchives)302{303 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBinaryArchives_), binaryArchives);304}305306_MTL_INLINE void MTL::StitchedLibraryDescriptor::setFunctionGraphs(const NS::Array* functionGraphs)307{308 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setFunctionGraphs_), functionGraphs);309}310311_MTL_INLINE void MTL::StitchedLibraryDescriptor::setFunctions(const NS::Array* functions)312{313 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setFunctions_), functions);314}315316_MTL_INLINE void MTL::StitchedLibraryDescriptor::setOptions(MTL::StitchedLibraryOptions options)317{318 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setOptions_), options);319}320