Program Listing for File matx_allocator.hpp

            

            
/*
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* 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.
*/

#ifndef HOLOSCAN_UTILS_MATX_ALLOCATOR_HPP
#define HOLOSCAN_UTILS_MATX_ALLOCATOR_HPP

#include <cuda_runtime_api.h>

#include <cstddef>
#include <cstdint>
#include <exception>
#include <memory>
#include <new>
#include <stdexcept>

#include "holoscan/core/resources/gxf/allocator.hpp"
#include "holoscan/core/resources/gxf/cuda_allocator.hpp"
#include "holoscan/logger/logger.hpp"

namespace holoscan {

class MatXAllocator {
 public:
  explicit MatXAllocator(Allocator* allocator,
                         MemoryStorageType storage_type = MemoryStorageType::kDevice,
                         cudaStream_t stream = nullptr)
      : allocator_(allocator),
        storage_type_(storage_type),
        stream_(stream),
        cuda_allocator_(dynamic_cast<CudaAllocator*>(allocator)) {
    if (!allocator_) {
      throw std::invalid_argument("MatXAllocator: allocator must not be null");
    }
    // Async allocation (CudaAllocator + stream) only supports device memory.
    // The GXF CudaAllocator::allocate_async_abi has no storage-type parameter
    // and always allocates device memory. Fail loudly to prevent silent misuse.
    if (cuda_allocator_ && stream_ && storage_type_ != MemoryStorageType::kDevice) {
      throw std::invalid_argument(
          "MatXAllocator: async allocation (CudaAllocator + stream) only "
          "supports kDevice storage type; use the synchronous path (no "
          "stream) for other memory types");
    }
  }

  MatXAllocator(Allocator* allocator, cudaStream_t stream)
      : MatXAllocator(allocator, MemoryStorageType::kDevice, stream) {}

  explicit MatXAllocator(const std::shared_ptr<Allocator>& allocator,
                         MemoryStorageType storage_type = MemoryStorageType::kDevice,
                         cudaStream_t stream = nullptr)
      : MatXAllocator(allocator.get(), storage_type, stream) {}

  MatXAllocator(const std::shared_ptr<Allocator>& allocator, cudaStream_t stream)
      : MatXAllocator(allocator.get(), MemoryStorageType::kDevice, stream) {}

  // Copyable and movable (lightweight, non-owning).
  MatXAllocator(const MatXAllocator&) = default;
  MatXAllocator& operator=(const MatXAllocator&) = default;
  MatXAllocator(MatXAllocator&&) = default;
  MatXAllocator& operator=(MatXAllocator&&) = default;

  [[nodiscard]] void* allocate(size_t size) {
    if (size == 0) {
      return nullptr;
    }

    void* ptr = nullptr;
    if (cuda_allocator_ && stream_) {
      // Path 1: CudaAllocator with stream — use async allocation.
      ptr =
          static_cast<void*>(cuda_allocator_->allocate_async(static_cast<uint64_t>(size), stream_));
    } else {
      // Path 2: Synchronous allocation with explicit memory type.
      ptr = static_cast<void*>(allocator_->allocate(static_cast<uint64_t>(size), storage_type_));
    }
    if (!ptr) {
      HOLOSCAN_LOG_ERROR("MatXAllocator: failed to allocate {} bytes", size);
      throw std::bad_alloc();
    }
    return ptr;
  }

  void deallocate(void* ptr, [[maybe_unused]] size_t size) noexcept {
    if (!ptr) {
      return;
    }

    try {
      if (cuda_allocator_ && stream_) {
        // Path 1: CudaAllocator with stream — use GXF-level free_async for
        // status. The Holoscan CudaAllocator::free_async wrapper is void and
        // silently swallows errors. Call the GXF API directly to get
        // Expected<void> for error detection.
        auto* gxf_cuda_alloc = cuda_allocator_->get();
        if (gxf_cuda_alloc) {
          auto result = gxf_cuda_alloc->free_async(static_cast<nvidia::byte*>(ptr), stream_);
          if (result) {
            return;
          }
          HOLOSCAN_LOG_WARN(
              "MatXAllocator: free_async failed, "
              "synchronizing stream before sync free fallback");
        } else {
          HOLOSCAN_LOG_WARN(
              "MatXAllocator: GXF CudaAllocator unavailable, "
              "synchronizing stream before sync free fallback");
        }
        // Synchronize the stream before falling back to synchronous free
        // to prevent use-after-free if GPU operations are still in flight.
        cudaStreamSynchronize(stream_);
        allocator_->free(static_cast<nvidia::byte*>(ptr));
      } else if (stream_) {
        // Path 2: Non-CudaAllocator with stream (e.g., BlockMemoryPool).
        // Access GXF-level allocator for stream-aware free(ptr, stream).
        auto* gxf_allocator = allocator_->get();
        if (gxf_allocator) {
          auto maybe_result =
              gxf_allocator->free(static_cast<nvidia::byte*>(ptr), static_cast<void*>(stream_));
          if (maybe_result) {
            return;
          }
          HOLOSCAN_LOG_WARN(
              "MatXAllocator: stream-aware free failed, "
              "synchronizing stream before sync free fallback");
        } else {
          HOLOSCAN_LOG_WARN(
              "MatXAllocator: GXF allocator unavailable, "
              "synchronizing stream before sync free fallback");
        }
        // Synchronize the stream before falling back to synchronous free
        // to prevent use-after-free if GPU operations are still in flight.
        cudaStreamSynchronize(stream_);
        allocator_->free(static_cast<nvidia::byte*>(ptr));
      } else {
        // Path 3: No stream — synchronous free.
        allocator_->free(static_cast<nvidia::byte*>(ptr));
      }
    } catch (const std::exception& e) {
      HOLOSCAN_LOG_ERROR("MatXAllocator: deallocate failed: {}", e.what());
    } catch (...) {
      HOLOSCAN_LOG_ERROR("MatXAllocator: deallocate failed with unknown exception");
    }
  }

  Allocator* allocator() const noexcept { return allocator_; }

  MemoryStorageType storage_type() const noexcept { return storage_type_; }

  cudaStream_t stream() const noexcept { return stream_; }

  MatXAllocator with_stream(cudaStream_t stream) const {
    return MatXAllocator(allocator_, storage_type_, stream);
  }

 private:
  Allocator* allocator_;
  MemoryStorageType storage_type_;
  cudaStream_t stream_;
  CudaAllocator* cuda_allocator_;
};

}  // namespace holoscan

#endif// HOLOSCAN_UTILS_MATX_ALLOCATOR_HPP