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