456 lines
12 KiB
C++
456 lines
12 KiB
C++
|
/*************************************************************************
|
||
|
|
* Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
|
||
|
|
*
|
||
|
|
* See LICENSE.txt for license information
|
||
|
|
************************************************************************/
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include "nccl.h"
|
||
|
|
#include "net.h"
|
||
|
|
#include "transport.h"
|
||
|
|
#include <cstdio>
|
||
|
|
#include <cstdlib>
|
||
|
|
#include <hip/hip_runtime.h>
|
||
|
|
#include <utility>
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @file ResourceGuards.hpp
|
||
|
|
* @brief Comprehensive RAII resource guards for automatic cleanup in tests
|
||
|
|
*
|
||
|
|
* Provides all RAII guard types for automatic resource management:
|
||
|
|
* - ScopeGuard: Generic cleanup for any action (with lambdas)
|
||
|
|
* - AutoGuard: Typed guards for resources with simple cleanup functions
|
||
|
|
* - ResourceGuard: Typed guards for resources with stateful deleters
|
||
|
|
* - Specialized guards: NcclRegHandleGuard, etc.
|
||
|
|
*
|
||
|
|
* Guards ensure cleanup even when ASSERT_* fails in tests.
|
||
|
|
* See MPITestRunner.md for detailed usage documentation.
|
||
|
|
*/
|
||
|
|
|
||
|
|
namespace RCCLTestGuards
|
||
|
|
{
|
||
|
|
|
||
|
|
// ============================================================================
|
||
|
|
// ScopeGuard - Generic cleanup for arbitrary actions
|
||
|
|
// ============================================================================
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @class ScopeGuard
|
||
|
|
* @brief Generic RAII scope guard for custom cleanup logic
|
||
|
|
*
|
||
|
|
* Executes a cleanup function on scope exit (normal return, early return, or exception).
|
||
|
|
* Useful for resources that don't have dedicated RAII guards or for one-off cleanup needs.
|
||
|
|
*
|
||
|
|
* @par Example:
|
||
|
|
* @code
|
||
|
|
* void* buffer = nullptr;
|
||
|
|
* hipMalloc(&buffer, size);
|
||
|
|
* auto guard = makeScopeGuard([&]() { if(buffer) hipFree(buffer); });
|
||
|
|
* // Automatic cleanup on scope exit
|
||
|
|
* @endcode
|
||
|
|
*
|
||
|
|
* @tparam Func Callable type (lambda, function pointer, functor)
|
||
|
|
*/
|
||
|
|
template<typename Func>
|
||
|
|
class ScopeGuard
|
||
|
|
{
|
||
|
|
Func cleanup_; ///< Cleanup function to execute on scope exit
|
||
|
|
bool dismissed_; ///< If true, skip cleanup (for ownership transfer)
|
||
|
|
|
||
|
|
public:
|
||
|
|
explicit ScopeGuard(Func f) noexcept : cleanup_(std::move(f)), dismissed_(false) {}
|
||
|
|
|
||
|
|
~ScopeGuard() noexcept
|
||
|
|
{
|
||
|
|
if(!dismissed_)
|
||
|
|
{
|
||
|
|
cleanup_();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void dismiss() noexcept { dismissed_ = true; }
|
||
|
|
void restore() noexcept { dismissed_ = false; }
|
||
|
|
|
||
|
|
ScopeGuard(ScopeGuard&& other) noexcept
|
||
|
|
: cleanup_(std::move(other.cleanup_)), dismissed_(other.dismissed_)
|
||
|
|
{
|
||
|
|
other.dismissed_ = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
ScopeGuard& operator=(ScopeGuard&& other) noexcept
|
||
|
|
{
|
||
|
|
if(this != &other)
|
||
|
|
{
|
||
|
|
if(!dismissed_)
|
||
|
|
{
|
||
|
|
cleanup_();
|
||
|
|
}
|
||
|
|
cleanup_ = std::move(other.cleanup_);
|
||
|
|
dismissed_ = other.dismissed_;
|
||
|
|
other.dismissed_ = true;
|
||
|
|
}
|
||
|
|
return *this;
|
||
|
|
}
|
||
|
|
|
||
|
|
ScopeGuard(const ScopeGuard&) = delete;
|
||
|
|
ScopeGuard& operator=(const ScopeGuard&) = delete;
|
||
|
|
};
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @brief Factory function to create ScopeGuard with type deduction
|
||
|
|
*
|
||
|
|
* @par Example:
|
||
|
|
* @code
|
||
|
|
* auto guard = makeScopeGuard([&]() { cleanup(); });
|
||
|
|
* @endcode
|
||
|
|
*/
|
||
|
|
template<typename Func>
|
||
|
|
ScopeGuard<Func> makeScopeGuard(Func f)
|
||
|
|
{
|
||
|
|
return ScopeGuard<Func>(std::move(f));
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @def SCOPE_EXIT
|
||
|
|
* @brief Convenience macro for creating anonymous scope guards
|
||
|
|
*
|
||
|
|
* @par Example:
|
||
|
|
* @code
|
||
|
|
* void* buffer = nullptr;
|
||
|
|
* hipMalloc(&buffer, size);
|
||
|
|
* SCOPE_EXIT(if(buffer) hipFree(buffer));
|
||
|
|
* @endcode
|
||
|
|
*/
|
||
|
|
#define SCOPE_EXIT_CONCAT_IMPL(a, b) a##b
|
||
|
|
#define SCOPE_EXIT_CONCAT(a, b) SCOPE_EXIT_CONCAT_IMPL(a, b)
|
||
|
|
#define SCOPE_EXIT(code) \
|
||
|
|
auto SCOPE_EXIT_CONCAT(scope_guard_, __LINE__) = RCCLTestGuards::makeScopeGuard([&]() { code; })
|
||
|
|
|
||
|
|
// ============================================================================
|
||
|
|
// AutoGuard & ResourceGuard - Typed resource management
|
||
|
|
// ============================================================================
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @class AutoGuard
|
||
|
|
* @brief Modern RAII guard using non-type template parameter for deleter
|
||
|
|
*
|
||
|
|
* Uses C++17's auto template parameters to directly reference cleanup functions,
|
||
|
|
* eliminating the need for deleter functors in simple cases.
|
||
|
|
*
|
||
|
|
* @tparam T Resource handle type
|
||
|
|
* @tparam DeleterFunc Function pointer for cleanup (auto-deduced)
|
||
|
|
*/
|
||
|
|
template<typename T, auto DeleterFunc>
|
||
|
|
class AutoGuard
|
||
|
|
{
|
||
|
|
private:
|
||
|
|
T resource_;
|
||
|
|
bool dismissed_;
|
||
|
|
|
||
|
|
public:
|
||
|
|
explicit AutoGuard(T resource = T{}) : resource_(resource), dismissed_(false) {}
|
||
|
|
|
||
|
|
~AutoGuard()
|
||
|
|
{
|
||
|
|
if(!dismissed_ && resource_)
|
||
|
|
{
|
||
|
|
DeleterFunc(resource_);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get the resource handle
|
||
|
|
T get() const
|
||
|
|
{
|
||
|
|
return resource_;
|
||
|
|
}
|
||
|
|
// Get pointer to resource handle (for API calls)
|
||
|
|
T* ptr()
|
||
|
|
{
|
||
|
|
return &resource_;
|
||
|
|
}
|
||
|
|
// Set the resource handle
|
||
|
|
void set(T resource)
|
||
|
|
{
|
||
|
|
resource_ = resource;
|
||
|
|
}
|
||
|
|
// Dismiss the guard (prevent cleanup)
|
||
|
|
void dismiss()
|
||
|
|
{
|
||
|
|
dismissed_ = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Release ownership (prevent cleanup)
|
||
|
|
T release()
|
||
|
|
{
|
||
|
|
dismissed_ = true;
|
||
|
|
return resource_;
|
||
|
|
}
|
||
|
|
|
||
|
|
AutoGuard(const AutoGuard&) = delete;
|
||
|
|
AutoGuard& operator=(const AutoGuard&) = delete;
|
||
|
|
|
||
|
|
AutoGuard(AutoGuard&& other) noexcept : resource_(other.resource_), dismissed_(other.dismissed_)
|
||
|
|
{
|
||
|
|
other.dismissed_ = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
AutoGuard& operator=(AutoGuard&& other) noexcept
|
||
|
|
{
|
||
|
|
if(this != &other)
|
||
|
|
{
|
||
|
|
if(!dismissed_ && resource_)
|
||
|
|
{
|
||
|
|
DeleterFunc(resource_);
|
||
|
|
}
|
||
|
|
resource_ = other.resource_;
|
||
|
|
dismissed_ = other.dismissed_;
|
||
|
|
other.dismissed_ = true;
|
||
|
|
}
|
||
|
|
return *this;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
/**
|
||
|
|
* @class ResourceGuard
|
||
|
|
* @brief Generic RAII guard template for resources with complex cleanup
|
||
|
|
*
|
||
|
|
* Uses a functor-based deleter for stateful deleters requiring additional context.
|
||
|
|
* For simple cleanup functions, prefer AutoGuard<T, func> instead.
|
||
|
|
*
|
||
|
|
* @tparam T Resource handle type
|
||
|
|
* @tparam Deleter Functor type for cleanup
|
||
|
|
*/
|
||
|
|
template<typename T, typename Deleter>
|
||
|
|
class ResourceGuard
|
||
|
|
{
|
||
|
|
private:
|
||
|
|
T resource_;
|
||
|
|
Deleter deleter_;
|
||
|
|
bool owns_;
|
||
|
|
|
||
|
|
public:
|
||
|
|
// Construct a resource guard
|
||
|
|
// @param resource Resource handle (can be nullptr/0)
|
||
|
|
// @param deleter Cleanup function/functor
|
||
|
|
explicit ResourceGuard(T resource = T{}, Deleter deleter = Deleter{})
|
||
|
|
: resource_(resource), deleter_(std::move(deleter)), owns_(true)
|
||
|
|
{}
|
||
|
|
|
||
|
|
// Destructor - automatically cleans up resource
|
||
|
|
~ResourceGuard()
|
||
|
|
{
|
||
|
|
if(owns_ && resource_)
|
||
|
|
{
|
||
|
|
deleter_(resource_);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Get the resource handle
|
||
|
|
T get() const
|
||
|
|
{
|
||
|
|
return resource_;
|
||
|
|
}
|
||
|
|
// Get pointer to resource handle (for API calls)
|
||
|
|
T* ptr()
|
||
|
|
{
|
||
|
|
return &resource_;
|
||
|
|
}
|
||
|
|
// Set the resource handle
|
||
|
|
void set(T resource)
|
||
|
|
{
|
||
|
|
resource_ = resource;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reset the resource handle
|
||
|
|
// @param resource New resource handle (can be nullptr/0)
|
||
|
|
void reset(T resource = T{})
|
||
|
|
{
|
||
|
|
if(owns_ && resource_ && resource_ != resource)
|
||
|
|
{
|
||
|
|
deleter_(resource_);
|
||
|
|
}
|
||
|
|
resource_ = resource;
|
||
|
|
owns_ = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
T release()
|
||
|
|
{
|
||
|
|
owns_ = false;
|
||
|
|
return resource_;
|
||
|
|
}
|
||
|
|
|
||
|
|
ResourceGuard(const ResourceGuard&) = delete;
|
||
|
|
ResourceGuard& operator=(const ResourceGuard&) = delete;
|
||
|
|
|
||
|
|
ResourceGuard(ResourceGuard&& other) noexcept
|
||
|
|
: resource_(other.resource_), deleter_(std::move(other.deleter_)), owns_(other.owns_)
|
||
|
|
{
|
||
|
|
other.owns_ = false;
|
||
|
|
}
|
||
|
|
|
||
|
|
ResourceGuard& operator=(ResourceGuard&& other) noexcept
|
||
|
|
{
|
||
|
|
if(this != &other)
|
||
|
|
{
|
||
|
|
// Clean up current resource
|
||
|
|
if(owns_ && resource_)
|
||
|
|
{
|
||
|
|
deleter_(resource_);
|
||
|
|
}
|
||
|
|
// Take ownership of other's resource
|
||
|
|
resource_ = other.resource_;
|
||
|
|
deleter_ = std::move(other.deleter_);
|
||
|
|
owns_ = other.owns_;
|
||
|
|
other.owns_ = false;
|
||
|
|
}
|
||
|
|
return *this;
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
// Note: Simple stateless deleters are replaced by wrapper functions + AutoGuard.
|
||
|
|
// Only stateful deleters that need additional context are kept here.
|
||
|
|
// Common deleters (NCCL-specific, used across many tests)
|
||
|
|
struct NcclRegHandleDeleter
|
||
|
|
{
|
||
|
|
ncclComm_t comm;
|
||
|
|
explicit NcclRegHandleDeleter(ncclComm_t c = nullptr) : comm(c) {}
|
||
|
|
void operator()(void* reg_handle) const
|
||
|
|
{
|
||
|
|
if(reg_handle && comm)
|
||
|
|
{
|
||
|
|
ncclCommDeregister(comm, reg_handle);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
// Wrapper functions for AutoGuard (void-returning cleanup functions)
|
||
|
|
inline void hipFreeWrapper(void* ptr)
|
||
|
|
{
|
||
|
|
if(ptr)
|
||
|
|
{
|
||
|
|
hipError_t err = hipFree(ptr);
|
||
|
|
if(err != hipSuccess)
|
||
|
|
{
|
||
|
|
fprintf(stderr,
|
||
|
|
"WARNING: hipFree failed in destructor: %s (ptr=%p)\n",
|
||
|
|
hipGetErrorString(err),
|
||
|
|
ptr);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline void hipStreamDestroyWrapper(hipStream_t stream)
|
||
|
|
{
|
||
|
|
if(stream)
|
||
|
|
{
|
||
|
|
hipError_t err = hipStreamDestroy(stream);
|
||
|
|
if(err != hipSuccess)
|
||
|
|
{
|
||
|
|
fprintf(stderr,
|
||
|
|
"WARNING: hipStreamDestroy failed in destructor: %s (stream=%p)\n",
|
||
|
|
hipGetErrorString(err),
|
||
|
|
static_cast<void*>(stream));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline void hipEventDestroyWrapper(hipEvent_t event)
|
||
|
|
{
|
||
|
|
if(event)
|
||
|
|
{
|
||
|
|
hipError_t err = hipEventDestroy(event);
|
||
|
|
if(err != hipSuccess)
|
||
|
|
{
|
||
|
|
fprintf(stderr,
|
||
|
|
"WARNING: hipEventDestroy failed in destructor: %s (event=%p)\n",
|
||
|
|
hipGetErrorString(err),
|
||
|
|
static_cast<void*>(event));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline void ncclCommDestroyWrapper(ncclComm_t comm)
|
||
|
|
{
|
||
|
|
if(comm)
|
||
|
|
{
|
||
|
|
ncclResult_t result = ncclCommDestroy(comm);
|
||
|
|
if(result != ncclSuccess)
|
||
|
|
{
|
||
|
|
fprintf(stderr,
|
||
|
|
"WARNING: ncclCommDestroy failed in destructor: %s (comm=%p)\n",
|
||
|
|
ncclGetErrorString(result),
|
||
|
|
static_cast<void*>(comm));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline void freeWrapper(void* ptr)
|
||
|
|
{
|
||
|
|
if(ptr)
|
||
|
|
free(ptr);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Type aliases for AutoGuard-based guards
|
||
|
|
using HostBufferAutoGuard = AutoGuard<void*, freeWrapper>;
|
||
|
|
using DeviceBufferAutoGuard = AutoGuard<void*, hipFreeWrapper>;
|
||
|
|
using HipStreamAutoGuard = AutoGuard<hipStream_t, hipStreamDestroyWrapper>;
|
||
|
|
using HipEventAutoGuard = AutoGuard<hipEvent_t, hipEventDestroyWrapper>;
|
||
|
|
using NcclCommAutoGuard = AutoGuard<ncclComm_t, ncclCommDestroyWrapper>;
|
||
|
|
|
||
|
|
// Type aliases for ResourceGuard-based guards (common/NCCL-specific)
|
||
|
|
using NcclRegHandleGuard = ResourceGuard<void*, NcclRegHandleDeleter>;
|
||
|
|
|
||
|
|
// Factory methods for ResourceGuard
|
||
|
|
template<typename T, typename Deleter>
|
||
|
|
inline auto makeGuard(T resource, Deleter deleter) -> ResourceGuard<T, Deleter>
|
||
|
|
{
|
||
|
|
return ResourceGuard<T, Deleter>(resource, std::move(deleter));
|
||
|
|
}
|
||
|
|
|
||
|
|
inline NcclRegHandleGuard makeRegHandleGuard(void* handle, ncclComm_t comm)
|
||
|
|
{
|
||
|
|
return NcclRegHandleGuard(handle, NcclRegHandleDeleter(comm));
|
||
|
|
}
|
||
|
|
|
||
|
|
template<typename T, typename Deleter>
|
||
|
|
inline auto makeCustomGuard(T resource, Deleter deleter) -> ResourceGuard<T, Deleter>
|
||
|
|
{
|
||
|
|
return ResourceGuard<T, Deleter>(resource, std::move(deleter));
|
||
|
|
}
|
||
|
|
|
||
|
|
// Factory methods for AutoGuard
|
||
|
|
template<typename T, auto DeleterFunc>
|
||
|
|
inline AutoGuard<T, DeleterFunc> makeAutoGuard(T resource)
|
||
|
|
{
|
||
|
|
return AutoGuard<T, DeleterFunc>(resource);
|
||
|
|
}
|
||
|
|
|
||
|
|
inline HostBufferAutoGuard makeHostBufferAutoGuard(void* buffer)
|
||
|
|
{
|
||
|
|
return HostBufferAutoGuard(buffer);
|
||
|
|
}
|
||
|
|
|
||
|
|
inline DeviceBufferAutoGuard makeDeviceBufferAutoGuard(void* buffer)
|
||
|
|
{
|
||
|
|
return DeviceBufferAutoGuard(buffer);
|
||
|
|
}
|
||
|
|
|
||
|
|
inline HipStreamAutoGuard makeStreamAutoGuard(hipStream_t stream)
|
||
|
|
{
|
||
|
|
return HipStreamAutoGuard(stream);
|
||
|
|
}
|
||
|
|
|
||
|
|
inline HipEventAutoGuard makeEventAutoGuard(hipEvent_t event)
|
||
|
|
{
|
||
|
|
return HipEventAutoGuard(event);
|
||
|
|
}
|
||
|
|
|
||
|
|
inline NcclCommAutoGuard makeCommAutoGuard(ncclComm_t comm)
|
||
|
|
{
|
||
|
|
return NcclCommAutoGuard(comm);
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace RCCLTestGuards
|
||
|
|
|