gpu_blas.cpp

Go to the documentation of this file.

// Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
// LICENSE and NOTICE for details. LLNL-CODE-806117.
//
// This file is part of the MFEM library. For more information and source code
// availability visit https://mfem.org.
//
// MFEM is free software; you can redistribute it and/or modify it under the
// terms of the BSD-3 license. We welcome feedback and contributions, see file
// CONTRIBUTING.md for details.
 
#include "gpu_blas.hpp"
#include "../../general/forall.hpp"
 
#if defined(MFEM_USE_CUDA)
#define MFEM_cu_or_hip(stub) cu##stub
#define MFEM_CU_or_HIP(stub) CU##stub
#elif defined(MFEM_USE_HIP)
#define MFEM_cu_or_hip(stub) hip##stub
#define MFEM_CU_or_HIP(stub) HIP##stub
#endif
 
#define MFEM_CONCAT(x, y, z) MFEM_CONCAT_(x, y, z)
#define MFEM_CONCAT_(x, y, z) x ## y ## z
 
#ifdef MFEM_USE_SINGLE
#define MFEM_GPUBLAS_PREFIX(stub) MFEM_CONCAT(MFEM_cu_or_hip(blas), S, stub)
#elif defined(MFEM_USE_DOUBLE)
#define MFEM_GPUBLAS_PREFIX(stub) MFEM_CONCAT(MFEM_cu_or_hip(blas), D, stub)
#endif
 
#define MFEM_BLAS_SUCCESS MFEM_CU_or_HIP(BLAS_STATUS_SUCCESS)
 
namespace mfem
{
 
GPUBlas &GPUBlas::Instance()
{
   static GPUBlas instance;
   return instance;
}
 
GPUBlas::HandleType GPUBlas::Handle()
{
   return Instance().handle;
}
 
#ifndef MFEM_USE_CUDA_OR_HIP
 
GPUBlas::GPUBlas() { }
GPUBlas::~GPUBlas() { }
void GPUBlas::EnableAtomics() { }
void GPUBlas::DisableAtomics() { }
 
#else
 
using blasStatus_t = MFEM_cu_or_hip(blasStatus_t);
 
GPUBlas::GPUBlas()
{
   blasStatus_t status = MFEM_cu_or_hip(blasCreate)(&handle);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "Cannot initialize GPU BLAS.");
}
 
GPUBlas::~GPUBlas()
{
   MFEM_cu_or_hip(blasDestroy)(handle);
}
 
void GPUBlas::EnableAtomics()
{
   const blasStatus_t status = MFEM_cu_or_hip(blasSetAtomicsMode)(
                                  Handle(), MFEM_CU_or_HIP(BLAS_ATOMICS_ALLOWED));
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "GPU BLAS error.");
}
 
void GPUBlas::DisableAtomics()
{
   const blasStatus_t status = MFEM_cu_or_hip(blasSetAtomicsMode)(
                                  Handle(), MFEM_CU_or_HIP(BLAS_ATOMICS_NOT_ALLOWED));
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "GPU BLAS error.");
}
 
void GPUBlasBatchedLinAlg::AddMult(const DenseTensor &A, const Vector &x,
                                   Vector &y, real_t alpha, real_t beta,
                                   Op op) const
{
   const bool tr = (op == Op::T);
 
   const int m = tr ? A.SizeJ() : A.SizeI();
   const int n = tr ? A.SizeI() : A.SizeJ();
   const int n_mat = A.SizeK();
   const int k = x.Size() / n / n_mat;
 
   auto d_A = A.Read();
   auto d_x = x.Read(); // Shape: (n, k, n_mat)
   auto d_y = beta == 0.0 ? y.Write() : y.ReadWrite(); // Shape (m, k, n_mat)
 
   const auto op_A = tr ? MFEM_CU_or_HIP(BLAS_OP_T) : MFEM_CU_or_HIP(BLAS_OP_N);
   const auto op_B = MFEM_CU_or_HIP(BLAS_OP_N);
 
   const blasStatus_t status = MFEM_GPUBLAS_PREFIX(gemmStridedBatched)(
                                  GPUBlas::Handle(), op_A, op_B, m, k, n,
                                  &alpha, d_A, m, m*n, d_x, n, n*k, &beta, d_y,
                                  m, m*k, n_mat);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "GPU BLAS error.");
}
 
void GPUBlasBatchedLinAlg::LUFactor(DenseTensor &A, Array<int> &P) const
{
   const int n = A.SizeI();
   const int n_mat = A.SizeK();
 
   P.SetSize(n*n_mat);
 
   Array<int> info_array(n_mat);
 
   real_t *A_base = A.ReadWrite();
   Array<real_t*> A_ptrs(n_mat);
   real_t **d_A_ptrs = A_ptrs.Write();
   mfem::forall(n_mat, [=] MFEM_HOST_DEVICE (int i)
   {
      d_A_ptrs[i] = A_base + i*n*n;
   });
 
   const blasStatus_t status = MFEM_GPUBLAS_PREFIX(getrfBatched)(
                                  GPUBlas::Handle(), n, d_A_ptrs, n, P.Write(),
                                  info_array.Write(), n_mat);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "");
}
 
void GPUBlasBatchedLinAlg::LUSolve(
   const DenseTensor &LU, const Array<int> &P, Vector &x) const
{
   const int n = LU.SizeI();
   const int n_mat = LU.SizeK();
   const int n_rhs = x.Size() / n / n_mat;
 
   Array<real_t*> A_ptrs(n_mat);
   real_t **d_A_ptrs = A_ptrs.Write();
   Array<real_t*> B_ptrs(n_mat);
   real_t **d_B_ptrs = B_ptrs.Write();
 
   {
      real_t *A_base = const_cast<real_t*>(LU.Read());
      real_t *B_base = x.ReadWrite();
      mfem::forall(n_mat, [=] MFEM_HOST_DEVICE (int i)
      {
         d_A_ptrs[i] = A_base + i*n*n;
         d_B_ptrs[i] = B_base + i*n*n_rhs;
      });
   }
 
   int info = 0;
   const blasStatus_t status = MFEM_GPUBLAS_PREFIX(getrsBatched)(
                                  GPUBlas::Handle(), MFEM_CU_or_HIP(BLAS_OP_N),
                                  n, n_rhs, d_A_ptrs, n, P.Read(), d_B_ptrs, n,
                                  &info, n_mat);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "");
}
 
void GPUBlasBatchedLinAlg::Invert(DenseTensor &A) const
{
   const int n = A.SizeI();
   const int n_mat = A.SizeK();
 
   DenseTensor LU(A.SizeI(), A.SizeJ(), A.SizeK());
   LU.Write();
   LU.GetMemory().CopyFrom(A.GetMemory(), A.TotalSize());
 
   Array<real_t*> LU_ptrs(n_mat);
   Array<real_t*> A_ptrs(n_mat);
   real_t **d_A_ptrs = A_ptrs.Write();
   real_t **d_LU_ptrs = LU_ptrs.Write();
   {
      real_t *A_base = A.ReadWrite();
      real_t *LU_base = LU.Write();
      mfem::forall(n_mat, [=] MFEM_HOST_DEVICE (int i)
      {
         d_A_ptrs[i] = A_base + i*n*n;
         d_LU_ptrs[i] = LU_base + i*n*n;
      });
   }
 
   Array<int> P(n*n_mat);
   Array<int> info_array(n_mat);
   blasStatus_t status;
 
   status = MFEM_GPUBLAS_PREFIX(getrfBatched)(
               GPUBlas::Handle(), n, d_LU_ptrs, n, P.Write(),
               info_array.Write(), n_mat);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "");
 
   status = MFEM_GPUBLAS_PREFIX(getriBatched)(
               GPUBlas::Handle(), n, d_LU_ptrs, n, P.ReadWrite(), d_A_ptrs, n,
               info_array.Write(), n_mat);
   MFEM_VERIFY(status == MFEM_BLAS_SUCCESS, "");
}
 
#endif
 
} // namespace mfem