text: /**
   * Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
   *
   * NVIDIA Corporation and its licensors retain all intellectual property
   * and proprietary rights in and to this software, related documentation
   * and any modifications thereto.  Any use, reproduction, disclosure or
   * distribution of this software and related documentation without an express
   * license agreement from NVIDIA Corporation is strictly prohibited.
   *
   */
  
  #include "cuda_runtime.h"
  #include "gstnvinfer_allocator.h"
  
  /* Standard GStreamer boiler plate macros */
  #define GST_TYPE_NVINFER_ALLOCATOR \
      (gst_nvinfer_allocator_get_type ())
  #define GST_NVINFER_ALLOCATOR(obj) \
      (G_TYPE_CHECK_INSTANCE_CAST((obj),GST_TYPE_NVINFER_ALLOCATOR,GstNvInferAllocator))
  #define GST_NVINFER_ALLOCATOR_CLASS(klass) \
      (G_TYPE_CHECK_CLASS_CAST((klass),GST_TYPE_NVINFER_ALLOCATOR,GstNvInferAllocatorClass))
  #define GST_IS_NVINFER_ALLOCATOR(obj) \
      (G_TYPE_CHECK_INSTANCE_TYPE((obj),GST_TYPE_NVINFER_ALLOCATOR))
  #define GST_IS_NVINFER_ALLOCATOR_CLASS(klass) \
      (G_TYPE_CHECK_CLASS_TYPE((klass),GST_TYPE_NVINFER_ALLOCATOR))
  
  typedef struct _GstNvInferAllocator GstNvInferAllocator;
  typedef struct _GstNvInferAllocatorClass GstNvInferAllocatorClass;
  
  G_GNUC_INTERNAL GType gst_nvinfer_allocator_get_type (void);
  
  GST_DEBUG_CATEGORY_STATIC (gst_nvinfer_allocator_debug);
  #define GST_CAT_DEFAULT gst_nvinfer_allocator_debug
  
  /**
   * Extends the GstAllocator class. Holds the parameters for allocator.
   */
  struct _GstNvInferAllocator
  {
    GstAllocator allocator;
    guint batch_size;
    guint width;
    guint height;
    NvBufSurfaceColorFormat color_format;
    guint gpu_id;
  };
  
  struct _GstNvInferAllocatorClass
  {
    GstAllocatorClass parent_class;
  };
  
  /* Standard boiler plate to create a debug category and initializing the
   * allocator type.
   */
  #define _do_init \
      GST_DEBUG_CATEGORY_INIT (gst_nvinfer_allocator_debug, "nvinferallocator", 0, "nvinfer allocator");
  #define gst_nvinfer_allocator_parent_class parent_class
  G_DEFINE_TYPE_WITH_CODE (GstNvInferAllocator, gst_nvinfer_allocator,
      GST_TYPE_ALLOCATOR, _do_init);
  
  /** Type of the memory allocated by the allocator. This can be used to identify
   * buffers / memories allocated by this allocator. */
  #define GST_NVINFER_MEMORY_TYPE "nvinfer"
  
  /** Structure allocated internally by the allocator. */
  typedef struct
  {
    /** Should be the first member of a structure extending GstMemory. */
    GstMemory mem;
    GstNvInferMemory mem_infer;
  } GstNvInferMem;
  
  /* Function called by GStreamer buffer pool to allocate memory using this
   * allocator. */
  static GstMemory *
  gst_nvinfer_allocator_alloc (GstAllocator * allocator, gsize size,
      GstAllocationParams * params)
  {
    GstNvInferAllocator *inferallocator = GST_NVINFER_ALLOCATOR (allocator);
    GstNvInferMem *nvmem = new GstNvInferMem;
    GstNvInferMemory *tmem = &nvmem->mem_infer;
    NvBufSurfaceCreateParams create_params = { 0 };
  
    create_params.gpuId = inferallocator->gpu_id;
    create_params.width = inferallocator->width;
    create_params.height = inferallocator->height;
    create_params.size = 0;
    create_params.isContiguous = 1;
    create_params.colorFormat = inferallocator->color_format;
    create_params.layout = NVBUF_LAYOUT_PITCH;
    create_params.memType = NVBUF_MEM_DEFAULT;
  
    if (NvBufSurfaceCreate (&tmem->surf, inferallocator->batch_size,
            &create_params) != 0) {
      GST_ERROR ("Error: Could not allocate internal buffer pool for nvinfer");
      return nullptr;
    }
  
    if(tmem->surf->memType == NVBUF_MEM_SURFACE_ARRAY) {
      if (NvBufSurfaceMapEglImage (tmem->surf, -1) != 0) {
        GST_ERROR ("Error: Could not map EglImage from NvBufSurface for nvinfer");
        return nullptr;
      }
  
      tmem->egl_frames.resize (inferallocator->batch_size);
      tmem->cuda_resources.resize (inferallocator->batch_size);
    }
  
    tmem->frame_memory_ptrs.assign (inferallocator->batch_size, nullptr);
  
    for (guint i = 0; i < inferallocator->batch_size; i++) {
      if(tmem->surf->memType == NVBUF_MEM_SURFACE_ARRAY) {
        if (cuGraphicsEGLRegisterImage (&tmem->cuda_resources[i],
                tmem->surf->surfaceList[i].mappedAddr.eglImage,
                CU_GRAPHICS_MAP_RESOURCE_FLAGS_NONE) != CUDA_SUCCESS) {
          g_printerr ("Failed to register EGLImage in cuda\n");
          return nullptr;
        }
  
        if (cuGraphicsResourceGetMappedEglFrame (&tmem->egl_frames[i],
                tmem->cuda_resources[i], 0, 0) != CUDA_SUCCESS) {
          g_printerr ("Failed to get mapped EGL Frame\n");
          return nullptr;
        }
        tmem->frame_memory_ptrs[i] = (char *) tmem->egl_frames[i].frame.pPitch[0];
      }
      else {
        /* Calculate pointers to individual frame memories in the batch memory and
        * insert in the vector. */
        tmem->frame_memory_ptrs[i] = (char *) tmem->surf->surfaceList[i].dataPtr;
      }
    }
  
    /* Initialize the GStreamer memory structure. */
    gst_memory_init ((GstMemory *) nvmem, (GstMemoryFlags) 0, allocator, nullptr,
        size, params->align, 0, size);
  
    return (GstMemory *) nvmem;
  }
  
  /* Function called by buffer pool for freeing memory using this allocator. */
  static void
  gst_nvinfer_allocator_free (GstAllocator * allocator, GstMemory * memory)
  {
    GstNvInferMem *nvmem = (GstNvInferMem *) memory;
    GstNvInferMemory *tmem = &nvmem->mem_infer;
    if(tmem->surf->memType == NVBUF_MEM_SURFACE_ARRAY) {
      GstNvInferAllocator *inferallocator = GST_NVINFER_ALLOCATOR (allocator);
  
      for (size_t i = 0; i < inferallocator->batch_size; i++) {
        cuGraphicsUnregisterResource (tmem->cuda_resources[i]);
      }
    }
  
    NvBufSurfaceUnMapEglImage (tmem->surf, -1);
    NvBufSurfaceDestroy (tmem->surf);
  
    delete nvmem;
  }
  
  /* Function called when mapping memory allocated by this allocator. Should
   * return pointer to GstNvInferMemory. */
  static gpointer
  gst_nvinfer_memory_map (GstMemory * mem, gsize maxsize, GstMapFlags flags)
  {
    GstNvInferMem *nvmem = (GstNvInferMem *) mem;
  
    return (gpointer) & nvmem->mem_infer;
  }
  
  static void
  gst_nvinfer_memory_unmap (GstMemory * mem)
  {
  }
  
  /* Standard boiler plate. Assigning implemented function pointers. */
  static void
  gst_nvinfer_allocator_class_init (GstNvInferAllocatorClass * klass)
  {
    GstAllocatorClass *allocator_class = GST_ALLOCATOR_CLASS (klass);
  
    allocator_class->alloc = GST_DEBUG_FUNCPTR (gst_nvinfer_allocator_alloc);
    allocator_class->free = GST_DEBUG_FUNCPTR (gst_nvinfer_allocator_free);
  }
  
  /* Standard boiler plate. Assigning implemented function pointers and setting
   * the memory type. */
  static void
  gst_nvinfer_allocator_init (GstNvInferAllocator * allocator)
  {
    GstAllocator *parent = GST_ALLOCATOR_CAST (allocator);
  
    parent->mem_type = GST_NVINFER_MEMORY_TYPE;
    parent->mem_map = gst_nvinfer_memory_map;
    parent->mem_unmap = gst_nvinfer_memory_unmap;
  }
  
  /* Create a new allocator of type GST_TYPE_NVINFER_ALLOCATOR and initialize
   * members. */
  GstAllocator *
  gst_nvinfer_allocator_new (guint width, guint height,
      NvBufSurfaceColorFormat color_format, guint batch_size, guint gpu_id)
  {
    GstNvInferAllocator *allocator = (GstNvInferAllocator *)
        g_object_new (GST_TYPE_NVINFER_ALLOCATOR,
        nullptr);
  
    allocator->width = width;
    allocator->height = height;
    allocator->batch_size = batch_size;
    allocator->gpu_id = gpu_id;
    allocator->color_format = color_format;
  
    return (GstAllocator *) allocator;
  }
  
  GstNvInferMemory *
  gst_nvinfer_buffer_get_memory (GstBuffer * buffer)
  {
    GstMemory *mem;
  
    mem = gst_buffer_peek_memory (buffer, 0);
  
    if (!mem || !gst_memory_is_type (mem, GST_NVINFER_MEMORY_TYPE))
      return nullptr;
  
    return &(((GstNvInferMem *) mem)->mem_infer);
  }
  