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/MTLArgumentEncoder.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"27#include "MTLDepthStencil.hpp"2829namespace MTL30{31class AccelerationStructure;32class ArgumentEncoder;33class Buffer;34class ComputePipelineState;35class Device;36class IndirectCommandBuffer;37class IntersectionFunctionTable;38class RenderPipelineState;39class SamplerState;40class Texture;41class VisibleFunctionTable;4243static const NS::UInteger AttributeStrideStatic = NS::UIntegerMax;4445class ArgumentEncoder : public NS::Referencing<ArgumentEncoder>46{47public:48 NS::UInteger alignment() const;4950 void* constantData(NS::UInteger index);5152 Device* device() const;5354 NS::UInteger encodedLength() const;5556 NS::String* label() const;5758 ArgumentEncoder* newArgumentEncoder(NS::UInteger index);5960 void setAccelerationStructure(const MTL::AccelerationStructure* accelerationStructure, NS::UInteger index);6162 void setArgumentBuffer(const MTL::Buffer* argumentBuffer, NS::UInteger offset);63 void setArgumentBuffer(const MTL::Buffer* argumentBuffer, NS::UInteger startOffset, NS::UInteger arrayElement);6465 void setBuffer(const MTL::Buffer* buffer, NS::UInteger offset, NS::UInteger index);66 void setBuffers(const MTL::Buffer* const buffers[], const NS::UInteger offsets[], NS::Range range);6768 void setComputePipelineState(const MTL::ComputePipelineState* pipeline, NS::UInteger index);69 void setComputePipelineStates(const MTL::ComputePipelineState* const pipelines[], NS::Range range);7071 void setDepthStencilState(const MTL::DepthStencilState* depthStencilState, NS::UInteger index);72 void setDepthStencilStates(const MTL::DepthStencilState* const depthStencilStates[], NS::Range range);7374 void setIndirectCommandBuffer(const MTL::IndirectCommandBuffer* indirectCommandBuffer, NS::UInteger index);75 void setIndirectCommandBuffers(const MTL::IndirectCommandBuffer* const buffers[], NS::Range range);7677 void setIntersectionFunctionTable(const MTL::IntersectionFunctionTable* intersectionFunctionTable, NS::UInteger index);78 void setIntersectionFunctionTables(const MTL::IntersectionFunctionTable* const intersectionFunctionTables[], NS::Range range);7980 void setLabel(const NS::String* label);8182 void setRenderPipelineState(const MTL::RenderPipelineState* pipeline, NS::UInteger index);83 void setRenderPipelineStates(const MTL::RenderPipelineState* const pipelines[], NS::Range range);8485 void setSamplerState(const MTL::SamplerState* sampler, NS::UInteger index);86 void setSamplerStates(const MTL::SamplerState* const samplers[], NS::Range range);8788 void setTexture(const MTL::Texture* texture, NS::UInteger index);89 void setTextures(const MTL::Texture* const textures[], NS::Range range);9091 void setVisibleFunctionTable(const MTL::VisibleFunctionTable* visibleFunctionTable, NS::UInteger index);92 void setVisibleFunctionTables(const MTL::VisibleFunctionTable* const visibleFunctionTables[], NS::Range range);93};9495}9697_MTL_INLINE NS::UInteger MTL::ArgumentEncoder::alignment() const98{99 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(alignment));100}101102_MTL_INLINE void* MTL::ArgumentEncoder::constantData(NS::UInteger index)103{104 return Object::sendMessage<void*>(this, _MTL_PRIVATE_SEL(constantDataAtIndex_), index);105}106107_MTL_INLINE MTL::Device* MTL::ArgumentEncoder::device() const108{109 return Object::sendMessage<MTL::Device*>(this, _MTL_PRIVATE_SEL(device));110}111112_MTL_INLINE NS::UInteger MTL::ArgumentEncoder::encodedLength() const113{114 return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(encodedLength));115}116117_MTL_INLINE NS::String* MTL::ArgumentEncoder::label() const118{119 return Object::sendMessage<NS::String*>(this, _MTL_PRIVATE_SEL(label));120}121122_MTL_INLINE MTL::ArgumentEncoder* MTL::ArgumentEncoder::newArgumentEncoder(NS::UInteger index)123{124 return Object::sendMessage<MTL::ArgumentEncoder*>(this, _MTL_PRIVATE_SEL(newArgumentEncoderForBufferAtIndex_), index);125}126127_MTL_INLINE void MTL::ArgumentEncoder::setAccelerationStructure(const MTL::AccelerationStructure* accelerationStructure, NS::UInteger index)128{129 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setAccelerationStructure_atIndex_), accelerationStructure, index);130}131132_MTL_INLINE void MTL::ArgumentEncoder::setArgumentBuffer(const MTL::Buffer* argumentBuffer, NS::UInteger offset)133{134 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setArgumentBuffer_offset_), argumentBuffer, offset);135}136137_MTL_INLINE void MTL::ArgumentEncoder::setArgumentBuffer(const MTL::Buffer* argumentBuffer, NS::UInteger startOffset, NS::UInteger arrayElement)138{139 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setArgumentBuffer_startOffset_arrayElement_), argumentBuffer, startOffset, arrayElement);140}141142_MTL_INLINE void MTL::ArgumentEncoder::setBuffer(const MTL::Buffer* buffer, NS::UInteger offset, NS::UInteger index)143{144 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBuffer_offset_atIndex_), buffer, offset, index);145}146147_MTL_INLINE void MTL::ArgumentEncoder::setBuffers(const MTL::Buffer* const buffers[], const NS::UInteger offsets[], NS::Range range)148{149 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBuffers_offsets_withRange_), buffers, offsets, range);150}151152_MTL_INLINE void MTL::ArgumentEncoder::setComputePipelineState(const MTL::ComputePipelineState* pipeline, NS::UInteger index)153{154 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setComputePipelineState_atIndex_), pipeline, index);155}156157_MTL_INLINE void MTL::ArgumentEncoder::setComputePipelineStates(const MTL::ComputePipelineState* const pipelines[], NS::Range range)158{159 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setComputePipelineStates_withRange_), pipelines, range);160}161162_MTL_INLINE void MTL::ArgumentEncoder::setDepthStencilState(const MTL::DepthStencilState* depthStencilState, NS::UInteger index)163{164 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDepthStencilState_atIndex_), depthStencilState, index);165}166167_MTL_INLINE void MTL::ArgumentEncoder::setDepthStencilStates(const MTL::DepthStencilState* const depthStencilStates[], NS::Range range)168{169 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDepthStencilStates_withRange_), depthStencilStates, range);170}171172_MTL_INLINE void MTL::ArgumentEncoder::setIndirectCommandBuffer(const MTL::IndirectCommandBuffer* indirectCommandBuffer, NS::UInteger index)173{174 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setIndirectCommandBuffer_atIndex_), indirectCommandBuffer, index);175}176177_MTL_INLINE void MTL::ArgumentEncoder::setIndirectCommandBuffers(const MTL::IndirectCommandBuffer* const buffers[], NS::Range range)178{179 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setIndirectCommandBuffers_withRange_), buffers, range);180}181182_MTL_INLINE void MTL::ArgumentEncoder::setIntersectionFunctionTable(const MTL::IntersectionFunctionTable* intersectionFunctionTable, NS::UInteger index)183{184 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setIntersectionFunctionTable_atIndex_), intersectionFunctionTable, index);185}186187_MTL_INLINE void MTL::ArgumentEncoder::setIntersectionFunctionTables(const MTL::IntersectionFunctionTable* const intersectionFunctionTables[], NS::Range range)188{189 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setIntersectionFunctionTables_withRange_), intersectionFunctionTables, range);190}191192_MTL_INLINE void MTL::ArgumentEncoder::setLabel(const NS::String* label)193{194 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setLabel_), label);195}196197_MTL_INLINE void MTL::ArgumentEncoder::setRenderPipelineState(const MTL::RenderPipelineState* pipeline, NS::UInteger index)198{199 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setRenderPipelineState_atIndex_), pipeline, index);200}201202_MTL_INLINE void MTL::ArgumentEncoder::setRenderPipelineStates(const MTL::RenderPipelineState* const pipelines[], NS::Range range)203{204 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setRenderPipelineStates_withRange_), pipelines, range);205}206207_MTL_INLINE void MTL::ArgumentEncoder::setSamplerState(const MTL::SamplerState* sampler, NS::UInteger index)208{209 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSamplerState_atIndex_), sampler, index);210}211212_MTL_INLINE void MTL::ArgumentEncoder::setSamplerStates(const MTL::SamplerState* const samplers[], NS::Range range)213{214 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSamplerStates_withRange_), samplers, range);215}216217_MTL_INLINE void MTL::ArgumentEncoder::setTexture(const MTL::Texture* texture, NS::UInteger index)218{219 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setTexture_atIndex_), texture, index);220}221222_MTL_INLINE void MTL::ArgumentEncoder::setTextures(const MTL::Texture* const textures[], NS::Range range)223{224 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setTextures_withRange_), textures, range);225}226227_MTL_INLINE void MTL::ArgumentEncoder::setVisibleFunctionTable(const MTL::VisibleFunctionTable* visibleFunctionTable, NS::UInteger index)228{229 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setVisibleFunctionTable_atIndex_), visibleFunctionTable, index);230}231232_MTL_INLINE void MTL::ArgumentEncoder::setVisibleFunctionTables(const MTL::VisibleFunctionTable* const visibleFunctionTables[], NS::Range range)233{234 Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setVisibleFunctionTables_withRange_), visibleFunctionTables, range);235}236