vart_npu_tensor.hpp

Go to the documentation of this file.

/*
 * Copyright (C) 2025-2026 Advanced Micro Devices, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace vart {
 
// Forward declaration of the private NpuTensor implementation
class NpuTensorPriv;
 
enum class DataType {
  UNKNOWN,  
  BOOLEAN,  
  INT8,     
  UINT8,    
  INT16,    
  UINT16,   
  BF16,     
  FP16,     
  INT32,    
  UINT32,   
  FLOAT32,  
  INT64,    
  UINT64,   
};
 
enum class MemoryLayout {
  UNKNOWN,    
  NC,         
  NCH,        
  NHC,        
  NHW,        
  NWC,        
  NHWC,       
  NCHW,       
  NHWC4,      
  NHWC8,      
  NC4HW4,     
  NC8HW8,     
  HCWNC4,     
  HCWNC8,     
  HCWNC16,    
  NHW16C4WC,  
  NHW16WC4C,  
  GENERIC,    
};
 
enum class MemoryType {
  UNKNOWN,               
  XRT_BO,                
  DMA_FD,                
  USER_POINTER_CMA,      
  USER_POINTER_NON_CMA,  
};
 
enum class TensorDirection {
  INPUT,   
  OUTPUT,  
};
 
enum class TensorType {
  CPU,  
  HW,   
};
 
struct NpuTensorInfo {
  std::string name;
  DataType data_type = DataType::UNKNOWN;
  TensorDirection direction = TensorDirection::INPUT;
  TensorType tensor_type = TensorType::HW;
  MemoryLayout memory_layout = MemoryLayout::UNKNOWN;
  std::vector<uint32_t> memory_layout_order;
  size_t size = 0;
  size_t size_in_bytes = 0;
  std::vector<uint32_t> shape;
  std::vector<uint32_t> strides;
 
  void print() const;
};
 
class NpuTensor {
 public:
  NpuTensor(const NpuTensorInfo& info, void* buffer, MemoryType mem_type);
 
  NpuTensor(const NpuTensorInfo& info, const void* buffer, MemoryType mem_type);
 
  NpuTensor();
 
  void* get_buffer();
 
  const void* get_buffer() const;
 
  void* get_virtual_address();
 
  const void* get_virtual_address() const;
 
  uint64_t get_physical_address() const;
 
  const NpuTensorInfo& get_info() const;
 
  MemoryType get_memory_type() const;
 
  void sync_buffer() const;
 
  int export_buffer() const;
 
  void print_info() const;
 
 private:
  friend class NpuTensorPrivAccess;      // Allow controlled access to the private implementation
  std::shared_ptr<NpuTensorPriv> priv_;  // NpuTensor private implementation
};
 
}  // namespace vart