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