spb/forge Public MIT

1//-------------------------------------------------------------------------------------------------------------------------------------------------------------2//3// Metal/MTLTensor.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 "MTLResource.hpp"28#include "MTLTypes.hpp"2930namespace MTL31{32class Buffer;33class TensorAuxiliaryPlaneDescriptor;34class TensorAuxiliaryPlaneDescriptorMap;35class TensorBufferAttachments;36class TensorDescriptor;37class TensorExtents;3839_MTL_CONST(NS::ErrorDomain, TensorDomain);4041_MTL_ENUM(NS::Integer, TensorDataType) {42    TensorDataTypeNone = 0,43    TensorDataTypeFloat32 = 3,44    TensorDataTypeFloat16 = 16,45    TensorDataTypeBFloat16 = 121,46    TensorDataTypeInt8 = 45,47    TensorDataTypeUInt8 = 49,48    TensorDataTypeInt16 = 37,49    TensorDataTypeUInt16 = 41,50    TensorDataTypeInt32 = 29,51    TensorDataTypeUInt32 = 33,52    TensorDataTypeInt4 = 143,53    TensorDataTypeUInt4 = 144,54    TensorDataTypeMetalFloat8UE8M0 = 145,55    TensorDataTypeUInt2 = 149,56    TensorDataTypeInt2 = 150,57    TensorDataTypeMetalFloat8E5M2 = 141,58    TensorDataTypeMetalFloat8E4M3 = 142,59    TensorDataTypeMetalFloat4E2M1 = 148,60};6162_MTL_ENUM(NS::Integer, TensorError) {63    TensorErrorNone = 0,64    TensorErrorInternalError = 1,65    TensorErrorInvalidDescriptor = 2,66};6768_MTL_ENUM(NS::Integer, TensorPlaneType) {69    TensorPlaneTypeData = 0,70    TensorPlaneTypeScales = 1,71};7273_MTL_OPTIONS(NS::UInteger, TensorUsage) {74    TensorUsageCompute = 1,75    TensorUsageRender = 1 << 1,76    TensorUsageMachineLearning = 1 << 2,77};7879class TensorExtents : public NS::Copying<TensorExtents>80{81public:82    static TensorExtents* alloc();8384    NS::Integer           extentAtDimensionIndex(NS::UInteger dimensionIndex);8586    TensorExtents*        init();87    TensorExtents*        init(NS::UInteger rank, const NS::Integer* values);8889    NS::UInteger          rank() const;90};91class TensorAuxiliaryPlaneDescriptor : public NS::Copying<TensorAuxiliaryPlaneDescriptor>92{93public:94    static TensorAuxiliaryPlaneDescriptor* alloc();9596    TensorExtents*                         blockFactors() const;9798    TensorDataType                         dataType() const;99100    TensorAuxiliaryPlaneDescriptor*        init();101102    void                                   setBlockFactors(const MTL::TensorExtents* blockFactors);103104    void                                   setDataType(MTL::TensorDataType dataType);105};106class TensorAuxiliaryPlaneDescriptorMap : public NS::Copying<TensorAuxiliaryPlaneDescriptorMap>107{108public:109    static TensorAuxiliaryPlaneDescriptorMap* alloc();110111    TensorAuxiliaryPlaneDescriptor*           descriptor(MTL::TensorPlaneType plane);112113    TensorAuxiliaryPlaneDescriptorMap*        init();114115    void                                      reset();116117    void                                      setDescriptor(const MTL::TensorAuxiliaryPlaneDescriptor* descriptor, MTL::TensorPlaneType plane);118};119class TensorDescriptor : public NS::Copying<TensorDescriptor>120{121public:122    static TensorDescriptor*           alloc();123124    TensorAuxiliaryPlaneDescriptorMap* auxiliaryPlanes() const;125126    CPUCacheMode                       cpuCacheMode() const;127128    TensorDataType                     dataType() const;129130    TensorExtents*                     dimensions() const;131132    HazardTrackingMode                 hazardTrackingMode() const;133134    TensorDescriptor*                  init();135136    ResourceOptions                    resourceOptions() const;137138    void                               setAuxiliaryPlanes(const MTL::TensorAuxiliaryPlaneDescriptorMap* auxiliaryPlanes);139140    void                               setCpuCacheMode(MTL::CPUCacheMode cpuCacheMode);141142    void                               setDataType(MTL::TensorDataType dataType);143144    void                               setDimensions(const MTL::TensorExtents* dimensions);145146    void                               setHazardTrackingMode(MTL::HazardTrackingMode hazardTrackingMode);147148    void                               setResourceOptions(MTL::ResourceOptions resourceOptions);149150    void                               setStorageMode(MTL::StorageMode storageMode);151152    void                               setStrides(const MTL::TensorExtents* strides);153154    void                               setUsage(MTL::TensorUsage usage);155156    StorageMode                        storageMode() const;157158    TensorExtents*                     strides() const;159160    TensorUsage                        usage() const;161};162class TensorBufferAttachments : public NS::Copying<TensorBufferAttachments>163{164public:165    static TensorBufferAttachments* alloc();166167    Buffer*                         buffer(MTL::TensorPlaneType plane);168169    TensorBufferAttachments*        init();170171    NS::UInteger                    offset(MTL::TensorPlaneType plane);172173    void                            reset();174175    void                            setBuffer(const MTL::Buffer* buffer, NS::UInteger offset, MTL::TensorPlaneType plane);176};177class TensorAuxiliaryPlane : public NS::Referencing<TensorAuxiliaryPlane>178{179public:180    TensorExtents*  blockFactors() const;181182    Buffer*         buffer() const;183    NS::UInteger    bufferOffset() const;184185    TensorDataType  dataType() const;186187    TensorPlaneType planeType() const;188};189class Tensor : public NS::Referencing<Tensor, Resource>190{191public:192    NS::Array*     auxiliaryPlanes() const;193194    Buffer*        buffer() const;195    NS::UInteger   bufferOffset() const;196197    TensorDataType dataType() const;198199    TensorExtents* dimensions() const;200201    void           getBytes(void* bytes, const MTL::TensorExtents* strides, const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions);202    void           getBytes(void* bytes, const MTL::TensorExtents* strides, const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, MTL::TensorPlaneType plane);203204    ResourceID     gpuResourceID() const;205206    void           replaceSliceOrigin(const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, const void* bytes, const MTL::TensorExtents* strides);207    void           replaceSliceOrigin(const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, MTL::TensorPlaneType plane, const void* bytes, const MTL::TensorExtents* strides);208209    TensorExtents* strides() const;210211    TensorUsage    usage() const;212};213214}215216_MTL_PRIVATE_DEF_CONST(NS::ErrorDomain, TensorDomain);217218_MTL_INLINE MTL::TensorExtents* MTL::TensorExtents::alloc()219{220    return NS::Object::alloc<MTL::TensorExtents>(_MTL_PRIVATE_CLS(MTLTensorExtents));221}222223_MTL_INLINE NS::Integer MTL::TensorExtents::extentAtDimensionIndex(NS::UInteger dimensionIndex)224{225    return Object::sendMessage<NS::Integer>(this, _MTL_PRIVATE_SEL(extentAtDimensionIndex_), dimensionIndex);226}227228_MTL_INLINE MTL::TensorExtents* MTL::TensorExtents::init()229{230    return NS::Object::init<MTL::TensorExtents>();231}232233_MTL_INLINE MTL::TensorExtents* MTL::TensorExtents::init(NS::UInteger rank, const NS::Integer* values)234{235    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(initWithRank_values_), rank, values);236}237238_MTL_INLINE NS::UInteger MTL::TensorExtents::rank() const239{240    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(rank));241}242243_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptor* MTL::TensorAuxiliaryPlaneDescriptor::alloc()244{245    return NS::Object::alloc<MTL::TensorAuxiliaryPlaneDescriptor>(_MTL_PRIVATE_CLS(MTLTensorAuxiliaryPlaneDescriptor));246}247248_MTL_INLINE MTL::TensorExtents* MTL::TensorAuxiliaryPlaneDescriptor::blockFactors() const249{250    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(blockFactors));251}252253_MTL_INLINE MTL::TensorDataType MTL::TensorAuxiliaryPlaneDescriptor::dataType() const254{255    return Object::sendMessage<MTL::TensorDataType>(this, _MTL_PRIVATE_SEL(dataType));256}257258_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptor* MTL::TensorAuxiliaryPlaneDescriptor::init()259{260    return NS::Object::init<MTL::TensorAuxiliaryPlaneDescriptor>();261}262263_MTL_INLINE void MTL::TensorAuxiliaryPlaneDescriptor::setBlockFactors(const MTL::TensorExtents* blockFactors)264{265    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBlockFactors_), blockFactors);266}267268_MTL_INLINE void MTL::TensorAuxiliaryPlaneDescriptor::setDataType(MTL::TensorDataType dataType)269{270    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDataType_), dataType);271}272273_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptorMap* MTL::TensorAuxiliaryPlaneDescriptorMap::alloc()274{275    return NS::Object::alloc<MTL::TensorAuxiliaryPlaneDescriptorMap>(_MTL_PRIVATE_CLS(MTLTensorAuxiliaryPlaneDescriptorMap));276}277278_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptor* MTL::TensorAuxiliaryPlaneDescriptorMap::descriptor(MTL::TensorPlaneType plane)279{280    return Object::sendMessage<MTL::TensorAuxiliaryPlaneDescriptor*>(this, _MTL_PRIVATE_SEL(descriptorForPlane_), plane);281}282283_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptorMap* MTL::TensorAuxiliaryPlaneDescriptorMap::init()284{285    return NS::Object::init<MTL::TensorAuxiliaryPlaneDescriptorMap>();286}287288_MTL_INLINE void MTL::TensorAuxiliaryPlaneDescriptorMap::reset()289{290    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(reset));291}292293_MTL_INLINE void MTL::TensorAuxiliaryPlaneDescriptorMap::setDescriptor(const MTL::TensorAuxiliaryPlaneDescriptor* descriptor, MTL::TensorPlaneType plane)294{295    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDescriptor_forPlane_), descriptor, plane);296}297298_MTL_INLINE MTL::TensorDescriptor* MTL::TensorDescriptor::alloc()299{300    return NS::Object::alloc<MTL::TensorDescriptor>(_MTL_PRIVATE_CLS(MTLTensorDescriptor));301}302303_MTL_INLINE MTL::TensorAuxiliaryPlaneDescriptorMap* MTL::TensorDescriptor::auxiliaryPlanes() const304{305    return Object::sendMessage<MTL::TensorAuxiliaryPlaneDescriptorMap*>(this, _MTL_PRIVATE_SEL(auxiliaryPlanes));306}307308_MTL_INLINE MTL::CPUCacheMode MTL::TensorDescriptor::cpuCacheMode() const309{310    return Object::sendMessage<MTL::CPUCacheMode>(this, _MTL_PRIVATE_SEL(cpuCacheMode));311}312313_MTL_INLINE MTL::TensorDataType MTL::TensorDescriptor::dataType() const314{315    return Object::sendMessage<MTL::TensorDataType>(this, _MTL_PRIVATE_SEL(dataType));316}317318_MTL_INLINE MTL::TensorExtents* MTL::TensorDescriptor::dimensions() const319{320    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(dimensions));321}322323_MTL_INLINE MTL::HazardTrackingMode MTL::TensorDescriptor::hazardTrackingMode() const324{325    return Object::sendMessage<MTL::HazardTrackingMode>(this, _MTL_PRIVATE_SEL(hazardTrackingMode));326}327328_MTL_INLINE MTL::TensorDescriptor* MTL::TensorDescriptor::init()329{330    return NS::Object::init<MTL::TensorDescriptor>();331}332333_MTL_INLINE MTL::ResourceOptions MTL::TensorDescriptor::resourceOptions() const334{335    return Object::sendMessage<MTL::ResourceOptions>(this, _MTL_PRIVATE_SEL(resourceOptions));336}337338_MTL_INLINE void MTL::TensorDescriptor::setAuxiliaryPlanes(const MTL::TensorAuxiliaryPlaneDescriptorMap* auxiliaryPlanes)339{340    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setAuxiliaryPlanes_), auxiliaryPlanes);341}342343_MTL_INLINE void MTL::TensorDescriptor::setCpuCacheMode(MTL::CPUCacheMode cpuCacheMode)344{345    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setCpuCacheMode_), cpuCacheMode);346}347348_MTL_INLINE void MTL::TensorDescriptor::setDataType(MTL::TensorDataType dataType)349{350    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDataType_), dataType);351}352353_MTL_INLINE void MTL::TensorDescriptor::setDimensions(const MTL::TensorExtents* dimensions)354{355    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setDimensions_), dimensions);356}357358_MTL_INLINE void MTL::TensorDescriptor::setHazardTrackingMode(MTL::HazardTrackingMode hazardTrackingMode)359{360    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setHazardTrackingMode_), hazardTrackingMode);361}362363_MTL_INLINE void MTL::TensorDescriptor::setResourceOptions(MTL::ResourceOptions resourceOptions)364{365    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setResourceOptions_), resourceOptions);366}367368_MTL_INLINE void MTL::TensorDescriptor::setStorageMode(MTL::StorageMode storageMode)369{370    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setStorageMode_), storageMode);371}372373_MTL_INLINE void MTL::TensorDescriptor::setStrides(const MTL::TensorExtents* strides)374{375    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setStrides_), strides);376}377378_MTL_INLINE void MTL::TensorDescriptor::setUsage(MTL::TensorUsage usage)379{380    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setUsage_), usage);381}382383_MTL_INLINE MTL::StorageMode MTL::TensorDescriptor::storageMode() const384{385    return Object::sendMessage<MTL::StorageMode>(this, _MTL_PRIVATE_SEL(storageMode));386}387388_MTL_INLINE MTL::TensorExtents* MTL::TensorDescriptor::strides() const389{390    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(strides));391}392393_MTL_INLINE MTL::TensorUsage MTL::TensorDescriptor::usage() const394{395    return Object::sendMessage<MTL::TensorUsage>(this, _MTL_PRIVATE_SEL(usage));396}397398_MTL_INLINE MTL::TensorBufferAttachments* MTL::TensorBufferAttachments::alloc()399{400    return NS::Object::alloc<MTL::TensorBufferAttachments>(_MTL_PRIVATE_CLS(MTLTensorBufferAttachments));401}402403_MTL_INLINE MTL::Buffer* MTL::TensorBufferAttachments::buffer(MTL::TensorPlaneType plane)404{405    return Object::sendMessage<MTL::Buffer*>(this, _MTL_PRIVATE_SEL(bufferForPlane_), plane);406}407408_MTL_INLINE MTL::TensorBufferAttachments* MTL::TensorBufferAttachments::init()409{410    return NS::Object::init<MTL::TensorBufferAttachments>();411}412413_MTL_INLINE NS::UInteger MTL::TensorBufferAttachments::offset(MTL::TensorPlaneType plane)414{415    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(offsetForPlane_), plane);416}417418_MTL_INLINE void MTL::TensorBufferAttachments::reset()419{420    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(reset));421}422423_MTL_INLINE void MTL::TensorBufferAttachments::setBuffer(const MTL::Buffer* buffer, NS::UInteger offset, MTL::TensorPlaneType plane)424{425    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setBuffer_offset_forPlane_), buffer, offset, plane);426}427428_MTL_INLINE MTL::TensorExtents* MTL::TensorAuxiliaryPlane::blockFactors() const429{430    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(blockFactors));431}432433_MTL_INLINE MTL::Buffer* MTL::TensorAuxiliaryPlane::buffer() const434{435    return Object::sendMessage<MTL::Buffer*>(this, _MTL_PRIVATE_SEL(buffer));436}437438_MTL_INLINE NS::UInteger MTL::TensorAuxiliaryPlane::bufferOffset() const439{440    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(bufferOffset));441}442443_MTL_INLINE MTL::TensorDataType MTL::TensorAuxiliaryPlane::dataType() const444{445    return Object::sendMessage<MTL::TensorDataType>(this, _MTL_PRIVATE_SEL(dataType));446}447448_MTL_INLINE MTL::TensorPlaneType MTL::TensorAuxiliaryPlane::planeType() const449{450    return Object::sendMessage<MTL::TensorPlaneType>(this, _MTL_PRIVATE_SEL(planeType));451}452453_MTL_INLINE NS::Array* MTL::Tensor::auxiliaryPlanes() const454{455    return Object::sendMessage<NS::Array*>(this, _MTL_PRIVATE_SEL(auxiliaryPlanes));456}457458_MTL_INLINE MTL::Buffer* MTL::Tensor::buffer() const459{460    return Object::sendMessage<MTL::Buffer*>(this, _MTL_PRIVATE_SEL(buffer));461}462463_MTL_INLINE NS::UInteger MTL::Tensor::bufferOffset() const464{465    return Object::sendMessage<NS::UInteger>(this, _MTL_PRIVATE_SEL(bufferOffset));466}467468_MTL_INLINE MTL::TensorDataType MTL::Tensor::dataType() const469{470    return Object::sendMessage<MTL::TensorDataType>(this, _MTL_PRIVATE_SEL(dataType));471}472473_MTL_INLINE MTL::TensorExtents* MTL::Tensor::dimensions() const474{475    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(dimensions));476}477478_MTL_INLINE void MTL::Tensor::getBytes(void* bytes, const MTL::TensorExtents* strides, const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions)479{480    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(getBytes_strides_fromSliceOrigin_sliceDimensions_), bytes, strides, sliceOrigin, sliceDimensions);481}482483_MTL_INLINE void MTL::Tensor::getBytes(void* bytes, const MTL::TensorExtents* strides, const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, MTL::TensorPlaneType plane)484{485    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(getBytes_strides_fromSliceOrigin_sliceDimensions_plane_), bytes, strides, sliceOrigin, sliceDimensions, plane);486}487488_MTL_INLINE MTL::ResourceID MTL::Tensor::gpuResourceID() const489{490    return Object::sendMessage<MTL::ResourceID>(this, _MTL_PRIVATE_SEL(gpuResourceID));491}492493_MTL_INLINE void MTL::Tensor::replaceSliceOrigin(const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, const void* bytes, const MTL::TensorExtents* strides)494{495    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(replaceSliceOrigin_sliceDimensions_withBytes_strides_), sliceOrigin, sliceDimensions, bytes, strides);496}497498_MTL_INLINE void MTL::Tensor::replaceSliceOrigin(const MTL::TensorExtents* sliceOrigin, const MTL::TensorExtents* sliceDimensions, MTL::TensorPlaneType plane, const void* bytes, const MTL::TensorExtents* strides)499{500    Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(replaceSliceOrigin_sliceDimensions_plane_withBytes_strides_), sliceOrigin, sliceDimensions, plane, bytes, strides);501}502503_MTL_INLINE MTL::TensorExtents* MTL::Tensor::strides() const504{505    return Object::sendMessage<MTL::TensorExtents*>(this, _MTL_PRIVATE_SEL(strides));506}507508_MTL_INLINE MTL::TensorUsage MTL::Tensor::usage() const509{510    return Object::sendMessage<MTL::TensorUsage>(this, _MTL_PRIVATE_SEL(usage));511}512