Skip to main content

7. MUSA数学加速库

在本章中，我们将给出几个最常用数学库的使用方法以及使用案例。MUSA数学库主要包含以下几部分：

  1. muBLAS(Basic Linear Algebra subroutines based on MUSA)：基于 MUSA开发的基本线性代数运算库。
  2. muFFT(Fast Fourier Transform based on MUSA) ：基于MUSA开发的快速傅里叶变换库。
  3. muRAND(pseudo-random and quasi-random numbers based on MUSA) ：基于MUSA随机数生成库。
  4. muSPARSE(basic linear algebra subroutines for sparse matrices based on MUSA): 基于MUSA 开发的稀疏矩阵数学库。
  5. muPP(Performance Primitives based on MUSA): 基于MUSA加速的图像和信号处理的函数库。

7.1. 线性代数库 muBLAS​

muBLAS是基于MUSA开发的基础线性代数库，在MTGPU上经过深度优化，在AI和HPC场景下被广泛使用。按照计算复杂性，muBLAS函数可分为三类，第一类用来处理标量、向量和向量与向量间的运算，第二类用来处理向量与矩阵之间的运算， 第三类用来进行矩阵与矩阵间的运算。下面代码展示muBLAS的使用

#include <cstdio>
#include <cstdlib>
#include <mublas.h>
#include <musa_runtime.h>
#include <vector>

int main(int argc, char* argv[])
{
    mublasHandle_t mublasH = NULL;
    musaStream_t   stream  = NULL;

    /*
     *   A = | 1.0 2.0 3.0 4.0 |
     *   B = | 5.0 6.0 7.0 8.0 |
     */

    const std::vector<float> A     = {1.0, 2.0, 3.0, 4.0};
    std::vector<float>       B     = {5.0, 6.0, 7.0, 8.0};
    const float              alpha = 2.1;
    const int                incx  = 1;
    const int                incy  = 1;

    float* d_A = nullptr;
    float* d_B = nullptr;

    /* step 1: create mublas handle, bind a stream */
    mublasCreate(&mublasH);

    musaStreamCreateWithFlags(&stream, musaStreamNonBlocking);
    mublasSetStream(mublasH, stream);

    /* step 2: copy data to device */
    musaMalloc(reinterpret_cast<void**>(&d_A), sizeof(float) * A.size());
    musaMalloc(reinterpret_cast<void**>(&d_B), sizeof(float) * B.size());

    musaMemcpyAsync(d_A, A.data(), sizeof(float) * A.size(), musaMemcpyHostToDevice, stream);
    musaMemcpyAsync(d_B, B.data(), sizeof(float) * B.size(), musaMemcpyHostToDevice, stream);
    /* step 3: compute */
    mublasSaxpy(mublasH, A.size(), &alpha, d_A, incx, d_B, incy);

    /* step 4: copy data to host */
    musaMemcpyAsync(B.data(), d_B, sizeof(float) * B.size(), musaMemcpyDeviceToHost, stream);
    musaStreamSynchronize(stream);

    /*
     *   B = | 7.10 10.20 13.30 16.40 |
     */
    printf("B\n");
    for(int i = 0; i < B.size(); i++)
        printf("%f ", B[i]);
    /* free resources */
    musaFree(d_A);
    musaFree(d_B);

    mublasDestroy(mublasH);

    musaStreamDestroy(stream);

    musaDeviceReset();

    return EXIT_SUCCESS;
}

7.2. 快速傅里叶变换库 muFFT​

muFFT 是基于MUSA开发的离散傅里叶变换的库，并在MTGPU上进行深度优化。在深度学习、计算机视觉、计算物理、分子动力学、量子化学、地震和医学成像等跨学科和商业利用使用广泛。下面代码展示muFFT的使用

#include <mufft.h>
#include <musa_runtime.h>

#include <complex>
#include <iostream>
#include <vector>

int main() {
  std::cout << "mufft 1D single-precision complex-to-complex transform\n";

  const int Nx = 8;
  std::vector<std::complex<float>> h_x(Nx);
  // Inititalize the input data
  for (size_t i = 0; i < Nx; i++) {
    h_x[i] = i;
  }
  std::cout << "input:\n";
  for (size_t i = 0; i < Nx; i++) {
    std::cout << h_x[i] << " ";
  }
  std::cout << std::endl;

  // Create MUSA device object and copy data to device
  size_t complex_bytes = sizeof(decltype(h_x)::value_type) * h_x.size();
  mufftComplex* d_x;
  musaMalloc(&d_x, complex_bytes);
  musaMemcpy(d_x, h_x.data(), complex_bytes, musaMemcpyHostToDevice);

  // Create the plan
  mufftHandle plan;
  mufftPlan1d(&plan,      // plan handle
              Nx,         // transform length
              MUFFT_C2C,  // transform type (MUFFT_C2C for single-precision)
              1);         // number of transforms
  // Execute plan:
  mufftExecC2C(plan, d_x, d_x, MUFFT_FORWARD);

  // Copy back the result to host
  musaMemcpy(h_x.data(), d_x, complex_bytes, musaMemcpyDeviceToHost);
  std::cout << "output:\n";
  for (size_t i = 0; i < Nx; i++) {
    std::cout << h_x[i] << " ";
  }
  std::cout << std::endl;

  // Release resources
  mufftDestroy(plan);
  musaFree(d_x);

  return 0;
}

7.3. 随机数生成库 muRAND​

muRAND是基于MUSA开发的高效伪随机数和准随机数生成器。下面代码展示muRAND的使用

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <musa_runtime.h>
#include <murand.h>
int main()
{
    int rand_n = 128;
    int seed = 777;

    musaStream_t stream;
    musaStreamCreateWithFlags(&stream, musaStreamNonBlocking);

    float *d_Rand;
    musaMalloc((void **)&d_Rand, rand_n * sizeof(float));

    murandGenerator_t prngGPU;
    murandCreateGenerator(&prngGPU, MURAND_RNG_PSEUDO_MTGP32);
    murandSetStream(prngGPU, stream);
    murandSetPseudoRandomGeneratorSeed(prngGPU, seed);

    float *h_RandGPU;
    musaMallocHost(&h_RandGPU, rand_n * sizeof(float));

    printf("Generating random numbers on GPU...\n\n");
    murandGenerateUniform(prngGPU, (float *) d_Rand, rand_n);

    printf("\nReading back the results...\n");
    musaMemcpyAsync(h_RandGPU, d_Rand, rand_n * sizeof(float), musaMemcpyDeviceToHost, stream);
    musaDeviceSynchronize();

    for(int i = 0; i < rand_n; i++)
    {
        printf("r[i] = %f \n", h_RandGPU[i]);
    }

    murandDestroyGenerator(prngGPU);
    musaStreamDestroy(stream);
    musaFree(d_Rand);
    musaFreeHost(h_RandGPU);

    exit(EXIT_SUCCESS);
    
}

7.4. 稀疏矩阵库 muSPARSE​

muSPARSE基于MUSA开发的稀疏矩阵运算库。在机器学习，计算流体动力学，地震勘探和计算科学等领域被广泛使用。下面代码展示muSPARSE的使用

#include <musa_runtime_api.h>
#include <iostream>
#include <musparse.h>

int main(int argc, char* argv[]) {
  // Query device
  int ndev;
  musaGetDeviceCount(&ndev);

  if (ndev < 1) {
    return -1;
  }
  // Query device properties
  musaDeviceProp prop;
  musaGetDeviceProperties(&prop, 0);
  std::cout << "Device: " << prop.name << std::endl;

  // muSPARSE handle
  musparseHandle_t handle;
  musparseCreateHandle(&handle);
  // Number of non-zeros of the sparse vector
  int nnz = 3;
  // Sparse index vector
  int hx_ind[3] = {0, 3, 5};
  // Sparse value vector
  float hx_val[3] = {1.0, 2.0, 3.0};
  // Dense vector
  float hy[9] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0};
  // Scalar alpha
  float alpha = 3.7;
  // Index base
  musparseIndexBase_t idx_base = MUSPARSE_INDEX_BASE_ZERO;
  // Offload data to device
  int* dx_ind;
  float* dx_val;
  float* dy;

  musaMalloc((void**)&dx_ind, sizeof(int) * nnz);
  musaMalloc((void**)&dx_val, sizeof(float) * nnz);
  musaMalloc((void**)&dy, sizeof(float) * 9);

  musaMemcpy(dx_ind, hx_ind, sizeof(int) * nnz, musaMemcpyHostToDevice);
  musaMemcpy(dx_val, hx_val, sizeof(float) * nnz, musaMemcpyHostToDevice);
  musaMemcpy(dy, hy, sizeof(float) * 9, musaMemcpyHostToDevice);

  // Call daxpyi to perform y = y + alpha * x
  musparseSaxpyi(handle, nnz, &alpha, dx_val, dx_ind, dy, idx_base);

  // Print result
  musaMemcpy(hy, dy, sizeof(float) * 9, musaMemcpyDeviceToHost);

  std::cout << "y:";

  for (int i = 0; i < 9; ++i) {
    std::cout << " " << hy[i];
  }
  std::cout << std::endl;
  // Clear muSPARSE
  musparseDestroyHandle(handle);

  // Clear device memory
  musaFree(dx_ind);
  musaFree(dx_val);
  musaFree(dy);
  return 0;
}

7.5. 二维图像与信号处理库 muPP​

muPP 是用于执行经 GPU 加速的图像、视频和信号处理的函数库，对标Intel Integrated Performance Primitives（IPP）和 NVIDIA Performance Primitives（NPP）。

muPP 库的主要功能侧重于图像处理，可以轻松执行颜色转换、图像压缩、过滤、阈值处理和几何变换等任务，广泛用于相关领域的开发工作，帮助开发人员在最大限度提高灵活性的同时保持函数的高性能。

通过简单的函数调用，muPP 库可以实现高性能的Low-Level图像或信号处理功能，帮助从事计算机视觉、工业检测、机器人、医学成像、通信、深度学习和高性能计算等领域的工程师和研究人员快速实现算法所需的相关功能。下面代码展示muPP的使用


#include <iostream>
#include <musa_runtime.h>
#include "muppdefs.h"
#include "muppi.h"

#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>

int main() {
  auto imgFile = "./Lena512.jpg";
  cv::Mat img = cv::imread(imgFile);

  MUppiSize roi{img.cols, img.rows};
  int dataSize = img.step[0] * img.rows;
  MUpp8u *muSrc, *muDst;
  int srcStep = img.step[0];
  int dstStep = img.step[0];

  musaMalloc(reinterpret_cast<void**>(&muSrc), dataSize);
  musaMalloc(reinterpret_cast<void**>(&muDst), dataSize);
  musaMemcpy(muSrc, img.data, dataSize, musaMemcpyHostToDevice);

  MUppStatus sts = MUppStatus::MUPP_SUCCESS;
  sts = muppiFilterGauss_8u_C3R(reinterpret_cast<const MUpp8u*>(muSrc), srcStep,
                                reinterpret_cast<MUpp8u*>(muDst), dstStep, roi,
                                MUPP_MASK_SIZE_15_X_15);

  if (MUppStatus::MUPP_SUCCESS != sts) {
    std::cerr << "MUPP executed abnormally!" << std::endl;
    std::cerr << "Please check input arguments!" << std::endl;
  } else {
    cv::Mat result = cv::Mat::zeros(img.rows, img.cols, CV_8UC3);
    musaMemcpy(result.data, muDst, dataSize, musaMemcpyDeviceToHost);
    cv::imwrite("result.jpg", result);
  }

  musaFree(muSrc);
  musaFree(muDst);
  return sts;
}

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。