Add AMD support via ort alternative to trt
This commit is contained in:
@@ -0,0 +1,177 @@
|
||||
// ── Gallery GEMM backend ──────────────────────────────────────────────────────
|
||||
// GPU similarity engine for the identity matcher. Uploads the reference gallery
|
||||
// once and computes the per-frame similarity matrix with a single SGEMM. The GPU
|
||||
// math library is selected at compile time by CMake (SAE_GEMM_BACKEND):
|
||||
// cuBLAS/CUDA or rocBLAS/HIP. This is the ONLY translation unit that includes
|
||||
// cublas/cuda or rocblas/hip headers.
|
||||
|
||||
#include "inference/similarity.hpp"
|
||||
|
||||
#if defined(SAE_GEMM_CUDA)
|
||||
#include <cublas_v2.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
#elif defined(SAE_GEMM_ROCM)
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <rocblas/rocblas.h>
|
||||
#else
|
||||
#error "gemm_backend.cpp requires SAE_GEMM_CUDA or SAE_GEMM_ROCM to be defined"
|
||||
#endif
|
||||
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
struct GpuError : std::runtime_error {
|
||||
using std::runtime_error::runtime_error;
|
||||
};
|
||||
|
||||
#if defined(SAE_GEMM_CUDA)
|
||||
|
||||
using stream_t = cudaStream_t;
|
||||
using blas_handle_t = cublasHandle_t;
|
||||
|
||||
inline void check_gpu(cudaError_t e, const char* what) {
|
||||
if (e != cudaSuccess)
|
||||
throw GpuError(std::string(what) + ": " + cudaGetErrorString(e));
|
||||
}
|
||||
inline void check_blas(cublasStatus_t s, const char* what) {
|
||||
if (s != CUBLAS_STATUS_SUCCESS)
|
||||
throw GpuError(std::string(what) + ": cublas error " + std::to_string(s));
|
||||
}
|
||||
|
||||
inline void gpu_malloc(void** p, size_t bytes) { check_gpu(cudaMalloc(p, bytes), "cudaMalloc"); }
|
||||
inline void gpu_free(void* p) { cudaFree(p); }
|
||||
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyHostToDevice, s), "H2D"); }
|
||||
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyDeviceToHost, s), "D2H"); }
|
||||
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n) { check_gpu(cudaMemcpy(dst, src, n, cudaMemcpyHostToDevice), "H2D_sync"); }
|
||||
inline void stream_create(stream_t* s) { check_gpu(cudaStreamCreate(s), "cudaStreamCreate"); }
|
||||
inline void stream_destroy(stream_t s) { cudaStreamDestroy(s); }
|
||||
inline void stream_sync(stream_t s) { check_gpu(cudaStreamSynchronize(s), "cudaStreamSync"); }
|
||||
inline void blas_create(blas_handle_t* h) { check_blas(cublasCreate(h), "cublasCreate"); }
|
||||
inline void blas_destroy(blas_handle_t h) { cublasDestroy(h); }
|
||||
inline void blas_set_stream(blas_handle_t h, stream_t s) { check_blas(cublasSetStream(h, s), "cublasSetStream"); }
|
||||
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
|
||||
const float* A, const float* B, float* C) {
|
||||
const float alpha = 1.f, beta = 0.f;
|
||||
check_blas(cublasSgemm(h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
m, n, k, &alpha, A, k, B, k, &beta, C, m),
|
||||
"cublasSgemm");
|
||||
}
|
||||
inline const char* backend_name() { return "cuBLAS/CUDA"; }
|
||||
|
||||
#else // SAE_GEMM_ROCM
|
||||
|
||||
using stream_t = hipStream_t;
|
||||
using blas_handle_t = rocblas_handle;
|
||||
|
||||
inline void check_gpu(hipError_t e, const char* what) {
|
||||
if (e != hipSuccess)
|
||||
throw GpuError(std::string(what) + ": " + hipGetErrorString(e));
|
||||
}
|
||||
inline void check_blas(rocblas_status s, const char* what) {
|
||||
if (s != rocblas_status_success)
|
||||
throw GpuError(std::string(what) + ": rocblas error " + std::to_string(s));
|
||||
}
|
||||
|
||||
inline void gpu_malloc(void** p, size_t bytes) { check_gpu(hipMalloc(p, bytes), "hipMalloc"); }
|
||||
inline void gpu_free(void* p) { (void)hipFree(p); }
|
||||
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyHostToDevice, s), "H2D"); }
|
||||
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyDeviceToHost, s), "D2H"); }
|
||||
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n) { check_gpu(hipMemcpy(dst, src, n, hipMemcpyHostToDevice), "H2D_sync"); }
|
||||
inline void stream_create(stream_t* s) { check_gpu(hipStreamCreate(s), "hipStreamCreate"); }
|
||||
inline void stream_destroy(stream_t s) { (void)hipStreamDestroy(s); }
|
||||
inline void stream_sync(stream_t s) { check_gpu(hipStreamSynchronize(s), "hipStreamSync"); }
|
||||
inline void blas_create(blas_handle_t* h) { check_blas(rocblas_create_handle(h), "rocblas_create_handle"); }
|
||||
inline void blas_destroy(blas_handle_t h) { rocblas_destroy_handle(h); }
|
||||
inline void blas_set_stream(blas_handle_t h, stream_t s) { check_blas(rocblas_set_stream(h, s), "rocblas_set_stream"); }
|
||||
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
|
||||
const float* A, const float* B, float* C) {
|
||||
const float alpha = 1.f, beta = 0.f;
|
||||
// rocblas_sgemm is column-major; same transposition trick as cuBLAS:
|
||||
// C(m×n) = A(k×m)^T * B(k×n) → S(N_gallery × n_faces) = G^T * Q
|
||||
check_blas(rocblas_sgemm(h, rocblas_operation_transpose, rocblas_operation_none,
|
||||
m, n, k, &alpha, A, k, B, k, &beta, C, m),
|
||||
"rocblas_sgemm");
|
||||
}
|
||||
inline const char* backend_name() { return "rocBLAS/HIP"; }
|
||||
|
||||
#endif
|
||||
|
||||
constexpr int kDim = 512;
|
||||
|
||||
class SimilarityEngine final : public ISimilarityEngine {
|
||||
public:
|
||||
SimilarityEngine(const float* gallery_row_major, int n_gallery, int max_faces)
|
||||
: n_gallery_(n_gallery), max_faces_(max_faces)
|
||||
{
|
||||
const size_t gallery_floats = static_cast<size_t>(n_gallery_) * kDim;
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_gallery_), gallery_floats * sizeof(float));
|
||||
gpu_memcpy_h2d_sync(d_gallery_, gallery_row_major, gallery_floats * sizeof(float));
|
||||
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_query_),
|
||||
static_cast<size_t>(max_faces_) * kDim * sizeof(float));
|
||||
gpu_malloc(reinterpret_cast<void**>(&d_sims_),
|
||||
static_cast<size_t>(max_faces_) * n_gallery_ * sizeof(float));
|
||||
|
||||
stream_create(&stream_);
|
||||
blas_create(&handle_);
|
||||
blas_set_stream(handle_, stream_);
|
||||
|
||||
host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
|
||||
|
||||
std::cerr << "[similarity] " << backend_name() << " engine: gallery resident on GPU ("
|
||||
<< (gallery_floats * sizeof(float)) / (1024 * 1024) << " MiB)\n";
|
||||
}
|
||||
|
||||
~SimilarityEngine() override {
|
||||
if (d_gallery_) gpu_free(d_gallery_);
|
||||
if (d_query_) gpu_free(d_query_);
|
||||
if (d_sims_) gpu_free(d_sims_);
|
||||
if (handle_) blas_destroy(handle_);
|
||||
if (stream_) stream_destroy(stream_);
|
||||
}
|
||||
|
||||
SimilarityEngine(const SimilarityEngine&) = delete;
|
||||
SimilarityEngine& operator=(const SimilarityEngine&) = delete;
|
||||
|
||||
int max_faces() const override { return max_faces_; }
|
||||
|
||||
const float* compute(const float* query_row_major, int n_faces) override {
|
||||
if (n_faces <= 0) return host_sims_.data();
|
||||
if (n_faces > max_faces_)
|
||||
throw std::runtime_error("SimilarityEngine: n_faces exceeds max_faces");
|
||||
|
||||
gpu_memcpy_h2d(d_query_, query_row_major,
|
||||
static_cast<size_t>(n_faces) * kDim * sizeof(float), stream_);
|
||||
|
||||
// S (N_gallery × n_faces) col-major = G(512 × N_gallery)^T * Q(512 × n_faces)
|
||||
blas_sgemm(handle_, n_gallery_, n_faces, kDim, d_gallery_, d_query_, d_sims_);
|
||||
|
||||
gpu_memcpy_d2h(host_sims_.data(), d_sims_,
|
||||
static_cast<size_t>(n_gallery_) * n_faces * sizeof(float), stream_);
|
||||
stream_sync(stream_);
|
||||
return host_sims_.data();
|
||||
}
|
||||
|
||||
private:
|
||||
int n_gallery_{0};
|
||||
int max_faces_{0};
|
||||
float* d_gallery_{nullptr};
|
||||
float* d_query_{nullptr};
|
||||
float* d_sims_{nullptr};
|
||||
std::vector<float> host_sims_;
|
||||
stream_t stream_{};
|
||||
blas_handle_t handle_{};
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<ISimilarityEngine> make_similarity_engine(
|
||||
const float* gallery_row_major, int n_gallery, int max_faces) {
|
||||
return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
|
||||
}
|
||||
Reference in New Issue
Block a user