NNETypes.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "UObject/Class.h"
 
#include "NNETypes.generated.h"
 
UENUM()
enum class ENNETensorDataType : uint8
{
    None,                               
    Char,                               
    Boolean,                            
    Half,                               
    Float,                              
    Double,                             
    Int8,                               
    Int16,                              
    Int32,                              
    Int64,                              
    UInt8,                              
    UInt16,                             
    UInt32,                             
    UInt64,                             
    Complex64,                          
    Complex128,                         
    BFloat16                            
};
 
namespace UE::NNE
{
    int32 NNE_API GetTensorDataTypeSizeInBytes(ENNETensorDataType InType);
 
    class FSymbolicTensorShape
    {
    public:
        constexpr static int32  MaxRank = 8;
 
    private:
        TArray<int32, TInlineAllocator<MaxRank>> Data;
 
    public:
        FSymbolicTensorShape() = default;
 
        FSymbolicTensorShape(const FSymbolicTensorShape& OtherShape) = default;
 
        static NNE_API FSymbolicTensorShape Make(TConstArrayView<int32> Data);
 
        inline TConstArrayView<int32> GetData() const { return Data; };
        
        inline int32 Rank() const { return Data.Num(); }
 
        NNE_API bool IsConcrete() const;
 
        NNE_API bool operator==(const FSymbolicTensorShape& OtherShape) const;
        
        NNE_API bool operator!=(const FSymbolicTensorShape& OtherShape) const;
 
        NNE_API void operator=(const FSymbolicTensorShape& OtherShape);
    };
 
    class FTensorShape
    {
    public:
        constexpr static int32  MaxRank = FSymbolicTensorShape::MaxRank;
 
    private:
        TArray<uint32, TInlineAllocator<MaxRank>> Data;
 
    public:
        FTensorShape() = default;
 
        FTensorShape(const FTensorShape& OtherShape) = default;
 
        static NNE_API FTensorShape Make(TConstArrayView<uint32> Data);
        
        static NNE_API FTensorShape MakeFromSymbolic(const FSymbolicTensorShape& SymbolicShape);
 
        inline TConstArrayView<uint32> GetData() const { return Data; };
        
        inline int32 Rank() const { return Data.Num(); }
        
        NNE_API uint64 Volume() const;
        
        NNE_API bool IsCompatibleWith(const FSymbolicTensorShape& SymbolicShape) const;
 
        NNE_API bool operator==(const FTensorShape& OtherShape) const;
        
        NNE_API bool operator!=(const FTensorShape& OtherShape) const;
 
        NNE_API void operator=(const FTensorShape& OtherShape);
    };
 
    class FTensorDesc
    {
    private:
        FString                 Name;
        ENNETensorDataType      DataType;
        FSymbolicTensorShape    Shape;
 
        FTensorDesc() = default;
 
    public:
        static NNE_API FTensorDesc Make(const FString& Name, const FSymbolicTensorShape& Shape, ENNETensorDataType DataType);
 
        inline const FString& GetName() const { return Name; }
        
        inline ENNETensorDataType GetDataType() const { return DataType; }
        
        inline uint32 GetElementByteSize() const { return GetTensorDataTypeSizeInBytes(DataType); }
 
        inline const FSymbolicTensorShape& GetShape() const { return Shape; }
    };
 
} // namespace UE::NNE