Add AMD support via ort alternative to trt

This commit is contained in:
2026-06-28 11:50:05 +02:00
parent a3ba53ddf7
commit 0ee131a692
27 changed files with 1357 additions and 977 deletions
+177
View File
@@ -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);
}