NNERuntimeRDG.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "Containers/ContainersFwd.h"
#include "NNEStatus.h"
#include "NNETypes.h"
#include "UObject/Interface.h"
 
#include "NNERuntimeRDG.generated.h"
 
class FRDGBuffer;
using FRDGBufferRef = FRDGBuffer*;
 
class FRDGBuilder;
 
class UNNEModelData;
 
namespace UE::NNE
{
 
struct FTensorBindingRDG
{
    FRDGBufferRef Buffer;
};
 
class IModelInstanceRDG
{
public:
 
    using ESetInputTensorShapesStatus = EResultStatus;
    using EEnqueueRDGStatus = EResultStatus;
 
    virtual ~IModelInstanceRDG() = default;
 
    virtual TConstArrayView<FTensorDesc> GetInputTensorDescs() const = 0;
 
    virtual TConstArrayView<FTensorDesc> GetOutputTensorDescs() const = 0;
 
    virtual TConstArrayView<FTensorShape> GetInputTensorShapes() const = 0;
 
    virtual TConstArrayView<FTensorShape> GetOutputTensorShapes() const = 0;
    
    virtual ESetInputTensorShapesStatus SetInputTensorShapes(TConstArrayView<FTensorShape> InInputShapes) = 0;
 
    virtual EEnqueueRDGStatus EnqueueRDG(FRDGBuilder& RDGBuilder, TConstArrayView<FTensorBindingRDG> Inputs, TConstArrayView<FTensorBindingRDG> Outputs) = 0;
};
 
class IModelRDG
{
public:
 
    virtual ~IModelRDG() = default;
 
    virtual TSharedPtr<UE::NNE::IModelInstanceRDG> CreateModelInstanceRDG() = 0;
};
 
} // UE::NNE
 
UINTERFACE(MinimalAPI)
class UNNERuntimeRDG : public UInterface
{
    GENERATED_BODY()
};
 
class INNERuntimeRDG
{
    GENERATED_BODY()
 
public:
 
    using ECanCreateModelRDGStatus = UE::NNE::EResultStatus;
 
    virtual ECanCreateModelRDGStatus CanCreateModelRDG(const TObjectPtr<UNNEModelData> ModelData) const = 0;
    
    virtual TSharedPtr<UE::NNE::IModelRDG> CreateModelRDG(const TObjectPtr<UNNEModelData> ModelData) = 0;
};