dedf08166a
Added 863 files from NVIDIA/cccl sparse checkout: - c2h/ (27 files): Catch2 test helpers — generators, validators, runner - nvbench_helper/ (10 files): Benchmark harness utilities - cmake/ (29 files): CMake presets and build helpers - cudax/ (794 files): Experimental CUDA extensions - AGENTS.md: NVIDIA's official AI agent instructions for CCCL - CMakePresets.json: Standardized build configurations - cccl-version.json: Version tracking Also added CCCL_ASSET_MAP.md mapping all 4295 CCCL files to competition value and PRD items. cccl_upstream now covers 100% of competition-critical assets: - 27 tuning headers (SM80/90/100 benchmark data) - 32 dispatch headers (algorithm implementations) - 60 Thrust examples (correctness verification) - 217 CUB Catch2 tests (regression matrix) - 153 CUB benchmarks (parameter space search) - 18 CUB examples (API verification) - 27 test helpers + benchmark harness - 794 cudax experimental extensions
1679 lines
45 KiB
Plaintext
1679 lines
45 KiB
Plaintext
//===----------------------------------------------------------------------===//
|
|
//
|
|
// Part of CUDASTF in CUDA C++ Core Libraries,
|
|
// under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/**
|
|
* @file
|
|
*
|
|
* @brief This example implements the POTRI matrix inversion algorithm over multiple devices
|
|
*
|
|
*
|
|
*/
|
|
#include <cuda/experimental/stf.cuh>
|
|
|
|
#include <iostream>
|
|
|
|
#include <cusolverDn.h>
|
|
|
|
#include <cublas_v2.h>
|
|
|
|
#define TILED
|
|
|
|
using namespace cuda::experimental::stf;
|
|
|
|
stream_ctx ctx;
|
|
|
|
static std::unordered_map<exec_place, cublasHandle_t, hash<exec_place>> cublas_handles;
|
|
static std::unordered_map<exec_place, cusolverDnHandle_t, hash<exec_place>> cusolver_handles;
|
|
|
|
/* Get a CUBLAS handle valid on the current execution place, or initialize it lazily */
|
|
cublasHandle_t& get_cublas_handle(const exec_place& ep = exec_place::current_device())
|
|
{
|
|
auto& result = cublas_handles[ep];
|
|
if (result == cublasHandle_t())
|
|
{ // not found, default value inserted
|
|
// Lazy initialization, and save the handle for future use
|
|
cuda_try(cublasCreate(&result));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
/* Get a CUSOLVER handle valid on the current execution place, or initialize it lazily */
|
|
cusolverDnHandle_t& get_cusolver_handle(const exec_place& ep = exec_place::current_device())
|
|
{
|
|
auto& result = cusolver_handles[ep];
|
|
if (result == cusolverDnHandle_t())
|
|
{ // not found, default value inserted
|
|
// Lazy initialization, and save the handle for future use
|
|
cuda_try(cusolverDnCreate(&result));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
template <typename T>
|
|
class matrix
|
|
{
|
|
public:
|
|
matrix(int NROWS, int NCOLS, int BLOCKSIZE_ROWS, int BLOCKSIZE_COLS, bool is_sym, const char* _symbol = "matrix")
|
|
{
|
|
symbol = _symbol;
|
|
|
|
sym_matrix = is_sym;
|
|
|
|
m = NROWS;
|
|
mb = BLOCKSIZE_ROWS;
|
|
|
|
n = NCOLS;
|
|
nb = BLOCKSIZE_COLS;
|
|
|
|
assert(m % mb == 0);
|
|
assert(n % nb == 0);
|
|
|
|
size_t s = m * n * sizeof(T);
|
|
// cuda_try(cudaMallocHost(&h_array, m*n*sizeof(T)));
|
|
// fprintf(stderr, "Allocating %ld x %ld x %ld = %ld bytes (%f GB) on host for %s\n", m, n, sizeof(T), s,
|
|
// s / (1024.0 * 1024.0 * 1024.0), _symbol);
|
|
h_array = (T*) malloc(s);
|
|
assert(h_array);
|
|
cuda_try(cudaHostRegister(h_array, s, cudaHostRegisterPortable));
|
|
// cuda_try(cudaMalloc(&d_array, m*n*sizeof(T)));
|
|
|
|
// Compute the number of blocks
|
|
mt = m / mb;
|
|
nt = n / nb;
|
|
|
|
handles.resize(mt * nt);
|
|
|
|
for (size_t colb = 0; colb < nt; colb++)
|
|
{
|
|
size_t low_rowb = sym_matrix ? colb : 0;
|
|
for (size_t rowb = low_rowb; rowb < mt; rowb++)
|
|
{
|
|
T* addr_h = get_block_h(rowb, colb);
|
|
auto& h = get_handle(rowb, colb);
|
|
|
|
#ifdef TILED
|
|
// tiles are stored contiguously
|
|
size_t ld = mb;
|
|
#else
|
|
size_t ld = m;
|
|
#endif
|
|
std::ignore = ld; // work around compiler bug
|
|
h = ctx.logical_data(make_slice(addr_h, std::tuple{mb, nb}, ld));
|
|
h.set_symbol(std::string(symbol) + "_" + std::to_string(rowb) + "_" + std::to_string(colb));
|
|
h.set_write_back(false);
|
|
}
|
|
}
|
|
|
|
cuda_try(cudaGetDeviceCount(&ndevs));
|
|
for (int a = 1; a * a <= ndevs; a++)
|
|
{
|
|
if (ndevs % a == 0)
|
|
{
|
|
grid_p = a;
|
|
grid_q = ndevs / a;
|
|
}
|
|
}
|
|
|
|
assert(grid_p * grid_q == ndevs);
|
|
|
|
// std::cout << "FOUND " << ndevs << " DEVICES "
|
|
// << "p=" << grid_p << " q=" << grid_q << '\n';
|
|
}
|
|
|
|
int get_preferred_devid(int row, int col)
|
|
{
|
|
return (row % grid_p) + (col % grid_q) * grid_p;
|
|
}
|
|
|
|
auto& get_handle(int row, int col)
|
|
{
|
|
return handles[row + col * mt];
|
|
}
|
|
|
|
size_t get_index(size_t row, size_t col)
|
|
{
|
|
#ifdef TILED
|
|
// Find which tile contains this element
|
|
int tile_row = row / mb;
|
|
int tile_col = col / nb;
|
|
|
|
size_t tile_size = mb * nb;
|
|
|
|
// Look for the index of the beginning of the tile
|
|
size_t tile_start = (tile_row + mt * tile_col) * tile_size;
|
|
|
|
// Offset within the tile
|
|
size_t offset = (row % mb) + (col % nb) * mb;
|
|
|
|
return tile_start + offset;
|
|
#else
|
|
return row + col * m;
|
|
#endif
|
|
}
|
|
|
|
T* get_block_h(int brow, int bcol)
|
|
{
|
|
size_t index = get_index(brow * mb, bcol * nb);
|
|
return &h_array[index];
|
|
}
|
|
|
|
// Fill with func(Matrix*,row, col)
|
|
template <typename Fun>
|
|
void fill(Fun&& fun)
|
|
{
|
|
nvtx_range r("fill");
|
|
|
|
// Fill blocks by blocks
|
|
for (size_t colb = 0; colb < nt; colb++)
|
|
{
|
|
size_t low_rowb = sym_matrix ? colb : 0;
|
|
for (size_t rowb = low_rowb; rowb < mt; rowb++)
|
|
{
|
|
// Each task fills a block
|
|
auto& h = get_handle(rowb, colb);
|
|
int devid = get_preferred_devid(rowb, colb);
|
|
|
|
ctx.parallel_for(exec_place::device(devid), h.shape(), h.write()).set_symbol("INIT")->*
|
|
[=] _CCCL_DEVICE(size_t lrow, size_t lcol, auto sA) {
|
|
size_t row = lrow + rowb * sA.extent(0);
|
|
size_t col = lcol + colb * sA.extent(1);
|
|
sA(lrow, lcol) = fun(row, col);
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
// Print blocks
|
|
void print()
|
|
{
|
|
// print blocks by blocks
|
|
for (size_t colb = 0; colb < nt; colb++)
|
|
{
|
|
int low_rowb = sym_matrix ? colb : 0;
|
|
for (size_t rowb = low_rowb; rowb < mt; rowb++)
|
|
{
|
|
// Each task fills a block
|
|
ctx.host_launch(get_handle(rowb, colb).read())->*[=](auto sA) {
|
|
for (size_t lcol = 0; lcol < sA.extent(1); lcol++)
|
|
{
|
|
size_t col = lcol + colb * sA.extent(1);
|
|
for (size_t lrow = 0; lrow < sA.extent(0); lrow++)
|
|
{
|
|
size_t row = lrow + rowb * sA.extent(0);
|
|
|
|
fprintf(stderr, "%d,%d : %le\n", row, col, sA(lrow, lcol));
|
|
}
|
|
}
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
T* h_array;
|
|
T* d_array;
|
|
size_t m; // nrows
|
|
size_t n; // ncols
|
|
|
|
// Is this a sym matrix ? (lower assumed)
|
|
bool sym_matrix;
|
|
|
|
size_t mb; // block size (rows)
|
|
size_t nb; // block size (cols)
|
|
|
|
size_t mt; // number of column blocks
|
|
size_t nt; // number of row blocks
|
|
|
|
// abstract data handles
|
|
std::vector<logical_data<slice<double, 2>>> handles;
|
|
|
|
const char* symbol;
|
|
|
|
// for the mapping
|
|
int ndevs;
|
|
int grid_p, grid_q;
|
|
};
|
|
|
|
void DPOTRF(cublasFillMode_t uplo, matrix<double>& A, int A_row, int A_col)
|
|
{
|
|
auto& Akk = A.get_handle(A_row, A_col);
|
|
size_t m_akk = Akk.shape().extent(0);
|
|
// Note that the handle may be different from the actual handle...
|
|
int Lwork_expected;
|
|
cuda_safe_call(cusolverDnDpotrf_bufferSize(get_cusolver_handle(), uplo, m_akk, nullptr, 0, &Lwork_expected));
|
|
|
|
auto potrf_buffer = ctx.logical_data(shape_of<slice<double>>(Lwork_expected));
|
|
potrf_buffer.set_allocator(ctx.get_default_allocator());
|
|
|
|
auto devInfo = ctx.logical_data(shape_of<slice<int>>(1));
|
|
|
|
auto t = ctx.task(Akk.rw(), potrf_buffer.write(), devInfo.write());
|
|
t.set_symbol("DPOTRF");
|
|
t->*[&](auto s, auto sAkk, auto buffer, auto info) {
|
|
auto& h = get_cusolver_handle();
|
|
cuda_try(cusolverDnSetStream(h, s));
|
|
|
|
cuda_try(cusolverDnDpotrf(
|
|
h,
|
|
uplo,
|
|
sAkk.extent(0),
|
|
sAkk.data_handle(),
|
|
sAkk.stride(1),
|
|
buffer.data_handle(),
|
|
buffer.extent(0),
|
|
info.data_handle()));
|
|
};
|
|
}
|
|
|
|
void DTRTRI(cublasFillMode_t uplo, cublasDiagType_t diag, matrix<double>& A, int A_row, int A_col)
|
|
{
|
|
// Preallocate a buffer used by CUSOLVER
|
|
size_t workspaceInBytesOnDevice, workspaceInBytesOnHost;
|
|
int64_t m_a00 = A.mb;
|
|
assert(A.mb == A.nb);
|
|
|
|
cuda_try(cusolverDnXtrtri_bufferSize(
|
|
get_cusolver_handle(),
|
|
uplo,
|
|
diag,
|
|
m_a00,
|
|
CUDA_R_64F /* DTRTRI */,
|
|
nullptr,
|
|
m_a00,
|
|
&workspaceInBytesOnDevice,
|
|
&workspaceInBytesOnHost));
|
|
|
|
// We don't support allocating buffers of 0 bytes ... XXX
|
|
if (workspaceInBytesOnHost == 0)
|
|
{
|
|
workspaceInBytesOnHost = 8;
|
|
}
|
|
|
|
auto d_buffer = ctx.logical_data(shape_of<slice<char>>(workspaceInBytesOnDevice));
|
|
auto h_buffer = ctx.logical_data(shape_of<slice<char>>(workspaceInBytesOnHost));
|
|
d_buffer.set_allocator(ctx.get_default_allocator());
|
|
h_buffer.set_allocator(ctx.get_default_allocator());
|
|
|
|
auto devInfo = ctx.logical_data(shape_of<slice<int>>(1));
|
|
|
|
auto t =
|
|
ctx.task(A.get_handle(A_row, A_col).rw(), d_buffer.write(), h_buffer.write(data_place::managed()), devInfo.write());
|
|
t.set_symbol("DTRTRI");
|
|
t->*[&](auto s, auto sA, auto dbuffer, auto hbuffer, auto info) {
|
|
auto& h = get_cusolver_handle();
|
|
cuda_try(cusolverDnSetStream(h, s));
|
|
|
|
// DTRTRI(...)
|
|
cuda_try(cusolverDnXtrtri(
|
|
h,
|
|
uplo,
|
|
diag,
|
|
sA.extent(0),
|
|
CUDA_R_64F /* DTRTRI */,
|
|
sA.data_handle(),
|
|
sA.stride(1),
|
|
(double*) dbuffer.data_handle(),
|
|
workspaceInBytesOnDevice,
|
|
(double*) hbuffer.data_handle(),
|
|
workspaceInBytesOnHost,
|
|
info.data_handle()));
|
|
};
|
|
}
|
|
|
|
/*
|
|
* Note: this code was taken from CUSOLVER
|
|
*
|
|
* SLACPY copies all or part of a two-dimensional matrix A to another matrix B.
|
|
*
|
|
* up up_and_lo
|
|
* 1 0 upper triangle, including diagonal
|
|
* 0 0 lower triangle, including diagonal
|
|
* ? 1 whole matrix
|
|
*
|
|
* configuration:
|
|
* dim3 grids( m/VEC, m/BY )
|
|
* dim3 threads(VEC,BY)
|
|
*/
|
|
template <typename T_ELEM_SRC, typename T_ELEM_DST, int VEC_LOG, int BY_LOG>
|
|
__global__ void __launch_bounds__(1 << (VEC_LOG + BY_LOG))
|
|
lacpy_kernel(int m, int n, const T_ELEM_SRC* A, size_t lda, T_ELEM_DST* B, size_t ldb, int up, int up_and_lo)
|
|
{
|
|
const int VEC = (1 << VEC_LOG);
|
|
const int BY = (1 << BY_LOG);
|
|
|
|
const int inx = threadIdx.x;
|
|
const int iny = threadIdx.y;
|
|
|
|
const int ibx = blockIdx.x * VEC;
|
|
const int iby = blockIdx.y * BY;
|
|
|
|
const int i = ibx + inx;
|
|
const int j = iby + iny;
|
|
|
|
if (ibx >= m)
|
|
{
|
|
return;
|
|
}
|
|
if (iby >= n)
|
|
{
|
|
return;
|
|
}
|
|
|
|
T_ELEM_SRC Areg = T_ELEM_SRC(0);
|
|
|
|
if (up_and_lo)
|
|
{
|
|
/*
|
|
* copy whole matrix
|
|
DO 60 J = 1, N
|
|
DO 50 I = 1, M
|
|
B( I, J ) = A( I, J )
|
|
50 CONTINUE
|
|
60 CONTINUE
|
|
*/
|
|
if ((i < m) && (j < n))
|
|
{
|
|
Areg = A[i + j * lda];
|
|
B[i + j * ldb] = T_ELEM_DST(Areg);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// only lower or upper triangle is copied.
|
|
if (up)
|
|
{
|
|
/*
|
|
* copy upper triangle, including diagonal
|
|
DO 20 J = 1, N
|
|
DO 10 I = 1, MIN( J, M )
|
|
B( I, J ) = A( I, J )
|
|
10 CONTINUE
|
|
20 CONTINUE
|
|
*/
|
|
if ((i <= min(j, m - 1)) && (j < n))
|
|
{
|
|
Areg = A[i + j * lda];
|
|
B[i + j * ldb] = T_ELEM_DST(Areg);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
/*
|
|
* copy lower triangle, including diagonal
|
|
DO 40 J = 1, N
|
|
DO 30 I = J, M
|
|
B( I, J ) = A( I, J )
|
|
30 CONTINUE
|
|
40 CONTINUE
|
|
|
|
*/
|
|
if (((j <= i) && (i < m)) && (j < n))
|
|
{
|
|
Areg = A[i + j * lda];
|
|
B[i + j * ldb] = T_ELEM_DST(Areg);
|
|
}
|
|
}
|
|
}
|
|
|
|
/*
|
|
* SLACPY copies all or part of a two-dimensional matrix A to another
|
|
* matrix B.
|
|
*
|
|
* Input
|
|
* -------
|
|
* UPLO is CHARACTER*1
|
|
* Specifies the part of the matrix A to be copied to B.
|
|
* = 'U': Upper triangular part
|
|
* = 'L': Lower triangular part
|
|
* Otherwise: All of the matrix A
|
|
*
|
|
* M is INTEGER
|
|
* The number of rows of the matrix A.
|
|
* M >= 0.
|
|
*
|
|
* N is INTEGER
|
|
* The number of columns of the matrix A.
|
|
* N >= 0.
|
|
*
|
|
* A is REAL array, dimension (LDA,N)
|
|
* The m by n matrix A. If UPLO = 'U', only the upper triangle
|
|
* or trapezoid is accessed; if UPLO = 'L', only the lower
|
|
* triangle or trapezoid is accessed.
|
|
*
|
|
* LDA is INTEGER
|
|
* The first dimension of the array A. LDA >= max(1,M).
|
|
*
|
|
* B is REAL array, dimension (LDB,N)
|
|
* On exit, B = A in the locations specified by UPLO.
|
|
*
|
|
* LDB is INTEGER
|
|
* The leading dimension of the array B. LDB >= max(1,M).
|
|
*
|
|
*/
|
|
template <typename T_ELEM_SRC, typename T_ELEM_DST>
|
|
cusolverStatus_t cusolverDnXlacpy(
|
|
cublasFillMode_t uplo, // "UPPER", B = upper(A)
|
|
// "LOWER", B = lower(A)
|
|
// otherwise, B = A
|
|
int m,
|
|
int n,
|
|
const T_ELEM_SRC* A,
|
|
int lda,
|
|
T_ELEM_DST* B,
|
|
int ldb,
|
|
cudaStream_t stream)
|
|
{
|
|
cusolverStatus_t status = CUSOLVER_STATUS_SUCCESS;
|
|
cudaError_t cudaStat1 = cudaSuccess;
|
|
|
|
int up = 0;
|
|
int up_and_lo = 0;
|
|
|
|
// Quick return if possible
|
|
if ((0 >= m) || (0 >= n))
|
|
{
|
|
return status;
|
|
}
|
|
|
|
/*
|
|
* up up_and_lo
|
|
* 1 0 upper triangle, including diagonal
|
|
* 0 0 lower triangle, including diagonal
|
|
* ? 1 whole matrix
|
|
*/
|
|
if (CUBLAS_FILL_MODE_LOWER == uplo)
|
|
{
|
|
// Lower triangular part
|
|
up = 0;
|
|
}
|
|
else if (CUBLAS_FILL_MODE_UPPER == uplo)
|
|
{
|
|
// upper triangular part
|
|
up = 1;
|
|
}
|
|
else
|
|
{
|
|
up_and_lo = 1; // Otherwise: All of the matrix A
|
|
}
|
|
|
|
const int VEC_LOG = 5;
|
|
const int BY_LOG = 3;
|
|
const int VEC = (1 << VEC_LOG);
|
|
const int BY = (1 << BY_LOG);
|
|
dim3 grids((m + VEC - 1) / VEC, (n + BY - 1) / BY);
|
|
dim3 threads(VEC, BY);
|
|
|
|
lacpy_kernel<T_ELEM_SRC, T_ELEM_DST, VEC_LOG, BY_LOG>
|
|
<<<grids, threads, 0, stream>>>(m, n, A, (size_t) lda, B, (size_t) ldb, up, up_and_lo);
|
|
|
|
cudaStat1 = cudaGetLastError(); /* launch error */
|
|
if (cudaSuccess != cudaStat1)
|
|
{
|
|
fprintf(stderr, "Error (lacpy): %d\n", cudaStat1);
|
|
status = CUSOLVER_STATUS_EXECUTION_FAILED;
|
|
}
|
|
|
|
return status;
|
|
}
|
|
|
|
cusolverStatus_t cusolverDnDlacpy(
|
|
cublasFillMode_t uplo, // "UPPER", B = upper(A)
|
|
// "LOWER", B = lower(A)
|
|
// otherwise, B = A
|
|
int m,
|
|
int n,
|
|
const double* A,
|
|
int lda,
|
|
double* B,
|
|
int ldb,
|
|
cudaStream_t stream)
|
|
{
|
|
return cusolverDnXlacpy<double, double>(uplo, m, n, A, lda, B, ldb, stream);
|
|
}
|
|
|
|
// Pretend there is a CUBLAS interface for DLAAUM
|
|
void cublasDnDlaaum_bufferSize(cublasHandle_t /*unused*/, int m, int n, size_t* Workspace_size)
|
|
{
|
|
assert(Workspace_size);
|
|
*Workspace_size = m * n * sizeof(double);
|
|
}
|
|
|
|
// Pretend there is a CUBLAS interface for DLAAUM
|
|
// A triangular
|
|
// Lower : A = A^T * A
|
|
// Upper : A = A A^T
|
|
void cublasDnDlaaum(
|
|
cublasHandle_t cublas_handle,
|
|
cublasFillMode_t uplo,
|
|
int m,
|
|
int n,
|
|
double* A,
|
|
int ldA,
|
|
double* Workspace_d,
|
|
size_t Workspace_size)
|
|
{
|
|
cudaStream_t stream;
|
|
cuda_safe_call(cublasGetStream(cublas_handle, &stream));
|
|
|
|
// "Hand coded"
|
|
// We use a full copy of A !
|
|
// fprintf(stderr, "GOT Workspace_size %ld ... expected %d\n", Workspace_size, m * n * sizeof(double));
|
|
std::ignore = Workspace_size;
|
|
assert(Workspace_size >= m * n * sizeof(double));
|
|
|
|
double* B = Workspace_d;
|
|
int ldB = m;
|
|
|
|
// Blank the buffer
|
|
cuda_safe_call(cudaMemsetAsync(B, 0, m * n * sizeof(double), stream));
|
|
|
|
// Copy A (with upper or lower 0 untouched)
|
|
cusolverDnDlacpy(uplo, m, n, A, ldA, B, ldB, stream);
|
|
|
|
cublasDiagType_t diag = CUBLAS_DIAG_NON_UNIT;
|
|
const double one = 1.0;
|
|
|
|
auto side = (uplo == CUBLAS_FILL_MODE_LOWER) ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
|
|
|
|
// LOWER: TRMM(A,B) : B = op(A) * B = A^T * B with A triangular (B = C in CUBLAS), CUBLAS_OP_T, CUBLAS_SIDE_RIGHT
|
|
// UPPER: TRMM(A,B) : B = B * op(A) = B A^T with A triangular (B = C in CUBLAS), CUBLAS_OP_T, CUBLAS_SIDE_RIGHT
|
|
cuda_safe_call(cublasDtrmm(cublas_handle, side, uplo, CUBLAS_OP_T, diag, m, n, &one, A, ldA, B, ldB, B, ldB));
|
|
|
|
// Copy B=AA^T back into A (with upper or lower 0 untouched)
|
|
cusolverDnDlacpy(uplo, m, n, B, ldB, A, ldA, stream);
|
|
}
|
|
|
|
void DLAAUM(cublasFillMode_t uplo, matrix<double>& A, int A_row, int A_col)
|
|
{
|
|
int NB = A.mb;
|
|
size_t Lwork;
|
|
cublasDnDlaaum_bufferSize(get_cublas_handle(), NB, NB, &Lwork);
|
|
|
|
auto d_buffer = ctx.logical_data(shape_of<slice<char>>(Lwork));
|
|
|
|
auto t = ctx.task(A.get_handle(A_row, A_col).rw(), d_buffer.write());
|
|
t.set_symbol("DLAAUM");
|
|
t->*[&](auto s, auto sA, auto buffer) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
cublasDnDlaaum(
|
|
h, uplo, sA.extent(0), sA.extent(1), sA.data_handle(), sA.stride(1), (double*) buffer.data_handle(), Lwork);
|
|
};
|
|
}
|
|
|
|
void DGEMM(
|
|
cublasOperation_t transa,
|
|
cublasOperation_t transb,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
int A_row,
|
|
int A_col,
|
|
matrix<double>& B,
|
|
int B_row,
|
|
int B_col,
|
|
double beta,
|
|
matrix<double>& C,
|
|
int C_row,
|
|
int C_col)
|
|
{
|
|
auto ignored = get_cublas_handle();
|
|
auto t =
|
|
ctx.task(A.get_handle(A_row, A_col).read(), B.get_handle(B_row, B_col).read(), C.get_handle(C_row, C_col).rw());
|
|
t.set_symbol("DGEMM");
|
|
t->*[&](auto s, auto sA, auto sB, auto sC) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
int k = (transa == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
|
|
cuda_try(cublasDgemm(
|
|
h,
|
|
transa,
|
|
transb,
|
|
sC.extent(0),
|
|
sC.extent(1),
|
|
k,
|
|
&alpha,
|
|
sA.data_handle(),
|
|
sA.stride(1),
|
|
sB.data_handle(),
|
|
sB.stride(1),
|
|
&beta,
|
|
sC.data_handle(),
|
|
sC.stride(1)));
|
|
};
|
|
}
|
|
|
|
void DSYMM(
|
|
cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
int A_row,
|
|
int A_col,
|
|
matrix<double>& B,
|
|
int B_row,
|
|
int B_col,
|
|
double beta,
|
|
matrix<double>& C,
|
|
int C_row,
|
|
int C_col)
|
|
{
|
|
auto ignored = get_cublas_handle();
|
|
auto t =
|
|
ctx.task(A.get_handle(A_row, A_col).read(), B.get_handle(B_row, B_col).read(), C.get_handle(C_row, C_col).rw());
|
|
t.set_symbol("DSYMM");
|
|
t->*[&](auto s, auto sA, auto sB, auto sC) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
cuda_try(cublasDsymm(
|
|
h,
|
|
side,
|
|
uplo,
|
|
sC.extent(0),
|
|
sC.extent(1),
|
|
&alpha,
|
|
sA.data_handle(),
|
|
sA.stride(1),
|
|
sB.data_handle(),
|
|
sB.stride(1),
|
|
&beta,
|
|
sC.data_handle(),
|
|
sC.stride(1)));
|
|
};
|
|
}
|
|
|
|
void DSYRK(
|
|
cublasFillMode_t uplo,
|
|
cublasOperation_t trans,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
int A_row,
|
|
int A_col,
|
|
double beta,
|
|
matrix<double>& C,
|
|
int C_row,
|
|
int C_col)
|
|
{
|
|
auto ignored = get_cublas_handle();
|
|
auto t = ctx.task(A.get_handle(A_row, A_col).read(), C.get_handle(C_row, C_col).rw());
|
|
t.set_symbol("DSYRK");
|
|
t->*[&](auto s, auto sA, auto sC) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
// number of rows of matrix op(A) and C
|
|
int n = sC.extent(0);
|
|
|
|
// number of columns of matrix op(A)
|
|
int k = (trans == CUBLAS_OP_N) ? sA.extent(1) : sA.extent(0);
|
|
|
|
cuda_try(
|
|
cublasDsyrk(h, uplo, trans, n, k, &alpha, sA.data_handle(), sA.stride(1), &beta, sC.data_handle(), sC.stride(1)));
|
|
};
|
|
}
|
|
|
|
void DTRSM(
|
|
cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
cublasOperation_t transa,
|
|
cublasDiagType_t diag,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
int A_row,
|
|
int A_col,
|
|
matrix<double>& B,
|
|
int B_row,
|
|
int B_col)
|
|
{
|
|
auto ignored = get_cublas_handle();
|
|
auto t = ctx.task(A.get_handle(A_row, A_col).read(), B.get_handle(B_row, B_col).rw());
|
|
t.set_symbol("DTRSM");
|
|
t->*[&](auto s, auto sA, auto sB) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
cuda_try(cublasDtrsm(
|
|
h,
|
|
side,
|
|
uplo,
|
|
transa,
|
|
diag,
|
|
sB.extent(0),
|
|
sB.extent(1),
|
|
&alpha,
|
|
sA.data_handle(),
|
|
sA.stride(1),
|
|
sB.data_handle(),
|
|
sB.stride(1)));
|
|
};
|
|
}
|
|
|
|
void DTRMM(
|
|
cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
cublasOperation_t transa,
|
|
cublasDiagType_t diag,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
int A_row,
|
|
int A_col,
|
|
matrix<double>& B,
|
|
int B_row,
|
|
int B_col)
|
|
{
|
|
auto ignored = get_cublas_handle();
|
|
auto t = ctx.task(A.get_handle(A_row, A_col).read(), B.get_handle(B_row, B_col).rw());
|
|
t.set_symbol("DTRMM");
|
|
t->*[&](auto s, auto sA, auto sB) {
|
|
auto& h = get_cublas_handle();
|
|
cuda_try(cublasSetStream(h, s));
|
|
|
|
// Note : CUBLAS DTRMM implementation is out of place but supports in place by using the same buffer B and C
|
|
cuda_try(cublasDtrmm(
|
|
get_cublas_handle(),
|
|
side,
|
|
uplo,
|
|
transa,
|
|
diag,
|
|
sB.extent(0),
|
|
sB.extent(1),
|
|
&alpha,
|
|
sA.data_handle(),
|
|
sA.stride(1),
|
|
sB.data_handle(),
|
|
sB.stride(1),
|
|
sB.data_handle(),
|
|
sB.stride(1) /* same as B*/));
|
|
};
|
|
}
|
|
|
|
void PDNRM2_HOST(matrix<double>* A, double* result)
|
|
{
|
|
#ifdef HAVE_DOT
|
|
ctx.get_dot()->set_current_color("red");
|
|
#endif
|
|
|
|
for (size_t rowb = 0; rowb < A->mt; rowb++)
|
|
{
|
|
for (size_t colb = 0; colb < A->nt; colb++)
|
|
{
|
|
ctx.host_launch(A->get_handle(rowb, colb).read())->*[=](auto sA) {
|
|
double res2 = 0.0;
|
|
for (size_t col = 0; col < sA.extent(1); col++)
|
|
{
|
|
for (size_t row = 0; row < sA.extent(0); row++)
|
|
{
|
|
double v = sA(row, col);
|
|
res2 += v * v;
|
|
}
|
|
}
|
|
*result += res2;
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
void PDPOTRF(matrix<double>& A)
|
|
{
|
|
nvtx_range r("PDPOTRF");
|
|
|
|
#ifdef HAVE_DOT
|
|
ctx.get_dot()->set_current_color("yellow");
|
|
#endif
|
|
|
|
assert(A.m == A.n);
|
|
assert(A.mt == A.nt);
|
|
|
|
int NBLOCKS = A.mt;
|
|
assert(A.mb == A.nb);
|
|
|
|
for (int K = 0; K < NBLOCKS; K++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(K, K)));
|
|
DPOTRF(CUBLAS_FILL_MODE_LOWER, A, K, K);
|
|
|
|
for (int row = K + 1; row < NBLOCKS; row++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(row, K)));
|
|
DTRSM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, K, K, A, row, K);
|
|
|
|
for (int col = K + 1; col < row; col++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(row, col)));
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, -1.0, A, row, K, A, col, K, 1.0, A, row, col);
|
|
}
|
|
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(row, row)));
|
|
DSYRK(CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, -1.0, A, row, K, 1.0, A, row, row);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Algorithm from PLASMA
|
|
void PDTRSM(cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
cublasOperation_t trans,
|
|
cublasDiagType_t diag,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
matrix<double>& B)
|
|
{
|
|
nvtx_range r("PDTRSM");
|
|
// std::cout << "[PDTRSM] START B MT " << B.mt << " NT " << B.nt << '\n';
|
|
|
|
if (side == CUBLAS_SIDE_LEFT)
|
|
{
|
|
if (uplo == CUBLAS_FILL_MODE_UPPER)
|
|
{
|
|
// TODO
|
|
abort();
|
|
}
|
|
else
|
|
{
|
|
//===========================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
|
|
//===========================================
|
|
if (trans == CUBLAS_OP_N)
|
|
{
|
|
for (size_t k = 0; k < B.mt; k++)
|
|
{
|
|
double lalpha = k == 0 ? alpha : 1.0;
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(k, k)));
|
|
DTRSM(side, uplo, trans, diag, lalpha, A, k, k, B, k, n);
|
|
}
|
|
for (size_t m = k + 1; m < B.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(m, k)));
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, -1.0, A, m, k, B, k, n, lalpha, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//================================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_[C|T]
|
|
//================================================
|
|
else
|
|
{
|
|
for (size_t k = 0; k < B.mt; k++)
|
|
{
|
|
double lalpha = k == 0 ? alpha : 1.0;
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - k - 1)));
|
|
DTRSM(side, uplo, trans, diag, lalpha, A, B.mt - k - 1, B.mt - k - 1, B, B.mt - k - 1, n);
|
|
}
|
|
for (size_t m = k + 1; m < B.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(B.mt - k - 1, B.mt - 1 - m)));
|
|
DGEMM(
|
|
trans, CUBLAS_OP_N, -1.0, A, B.mt - k - 1, B.mt - 1 - m, B, B.mt - k - 1, n, lalpha, B, B.mt - 1 - m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
// TODO
|
|
abort();
|
|
}
|
|
// std::cout << "[PDTRSM] END" << '\n';
|
|
}
|
|
|
|
void PDPOTRS(matrix<double>& A, matrix<double>& B, cublasFillMode_t uplo)
|
|
{
|
|
nvtx_range r("PDPOTRS");
|
|
|
|
#ifdef HAVE_DOT
|
|
ctx.get_dot()->set_current_color("green");
|
|
#endif
|
|
|
|
// std::cout << "[PDPOTRS] START" << '\n';
|
|
// Call the parallel functions.
|
|
PDTRSM(
|
|
CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_T : CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
|
|
|
|
#ifdef HAVE_DOT
|
|
ctx.get_dot()->set_current_color("darkgreen");
|
|
#endif
|
|
|
|
PDTRSM(
|
|
CUBLAS_SIDE_LEFT, uplo, uplo == CUBLAS_FILL_MODE_UPPER ? CUBLAS_OP_N : CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, B);
|
|
// std::cout << "[PDPOTRS] END" << '\n';
|
|
}
|
|
|
|
/***************************************************************************/ /**
|
|
* Parallel tile matrix-matrix
|
|
*multiplication.
|
|
* @see plasma_omp_dgemm
|
|
******************************************************************************/
|
|
void PDGEMM(cublasOperation_t transa,
|
|
cublasOperation_t transb,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
matrix<double>& B,
|
|
double beta,
|
|
matrix<double>& C)
|
|
{
|
|
nvtx_range r("PDGEMM");
|
|
|
|
#ifdef HAVE_DOT
|
|
reserved::dot::set_current_color("blue");
|
|
#endif
|
|
|
|
for (size_t m = 0; m < C.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < C.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(C.get_preferred_devid(m, n)));
|
|
|
|
//=========================================
|
|
// alpha*A*B does not contribute; scale C
|
|
//=========================================
|
|
int inner_k = transa == CUBLAS_OP_N ? A.n : A.m;
|
|
if (alpha == 0.0 || inner_k == 0)
|
|
{
|
|
DGEMM(transa, transb, alpha, A, 0, 0, B, 0, 0, beta, C, m, n);
|
|
}
|
|
else if (transa == CUBLAS_OP_N)
|
|
{
|
|
//================================
|
|
// CUBLAS_OP_N / CUBLAS_OP_N
|
|
//================================
|
|
if (transb == CUBLAS_OP_N)
|
|
{
|
|
assert(A.nt == B.mt);
|
|
for (size_t k = 0; k < A.nt; k++)
|
|
{
|
|
double zbeta = k == 0 ? beta : 1.0;
|
|
DGEMM(transa, transb, alpha, A, m, k, B, k, n, zbeta, C, m, n);
|
|
}
|
|
}
|
|
//=====================================
|
|
// CUBLAS_OP_N / CUBLAS_OP_T
|
|
//=====================================
|
|
else
|
|
{
|
|
for (size_t k = 0; k < A.nt; k++)
|
|
{
|
|
double zbeta = k == 0 ? beta : 1.0;
|
|
DGEMM(transa, transb, alpha, A, m, k, B, n, k, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
//=====================================
|
|
// CUBLAS_OP_T / CUBLAS_OP_N
|
|
//=====================================
|
|
if (transb == CUBLAS_OP_N)
|
|
{
|
|
for (size_t k = 0; k < A.mt; k++)
|
|
{
|
|
double zbeta = k == 0 ? beta : 1.0;
|
|
DGEMM(transa, transb, alpha, A, k, m, B, k, n, zbeta, C, m, n);
|
|
}
|
|
}
|
|
//==========================================
|
|
// CUBLAS_OP_T / CUBLAS_OP_T
|
|
//==========================================
|
|
else
|
|
{
|
|
for (size_t k = 0; k < A.mt; k++)
|
|
{
|
|
double zbeta = k == 0 ? beta : 1.0;
|
|
DGEMM(transa, transb, alpha, A, k, m, B, n, k, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Algorithm taken from the PLASMA library
|
|
*/
|
|
// We assume a lower triangular matrix (uplo == CUBLAS_FILL_MODE_LOWER)
|
|
void PDTRTRI(matrix<double>& A, cublasFillMode_t uplo, cublasDiagType_t diag)
|
|
{
|
|
nvtx_range r("PDTRTRI");
|
|
|
|
assert(uplo == CUBLAS_FILL_MODE_LOWER);
|
|
|
|
for (size_t k = 0; k < A.nt; k++)
|
|
{
|
|
for (size_t m = k + 1; m < A.mt; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(m, k)));
|
|
DTRSM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, diag, -1.0, A, k, k, A, m, k);
|
|
}
|
|
for (size_t m = k + 1; m < A.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < k; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(m, n)));
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, A, m, k, A, k, n, 1.0, A, m, n);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < k; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(k, n)));
|
|
DTRSM(CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, diag, 1.0, A, k, k, A, k, n);
|
|
}
|
|
|
|
// DTRTRI(...)
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(k, k)));
|
|
DTRTRI(uplo, diag, A, k, k);
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Algorithm taken from the PLASMA library
|
|
*/
|
|
// We assume a lower triangular matrix (uplo == CUBLAS_FILL_MODE_LOWER)
|
|
void PDLAUUM(matrix<double>& A, cublasFillMode_t uplo)
|
|
{
|
|
nvtx_range r("PDLAUUM");
|
|
assert(uplo == CUBLAS_FILL_MODE_LOWER);
|
|
|
|
for (size_t k = 0; k < A.mt; k++)
|
|
{
|
|
for (size_t n = 0; n < k; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(n, n)));
|
|
DSYRK(CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, 1.0, A, k, n, 1.0, A, n, n);
|
|
|
|
for (size_t m = n + 1; m < k; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(m, n)));
|
|
DGEMM(CUBLAS_OP_T, CUBLAS_OP_N, 1.0, A, k, m, A, k, n, 1.0, A, m, n);
|
|
}
|
|
}
|
|
for (size_t n = 0; n < k; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(k, n)));
|
|
DTRMM(CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, 1.0, A, k, k, A, k, n);
|
|
}
|
|
|
|
// LAAUM (Akk RW) (compute Akk^T * Akk)
|
|
cuda_try(cudaSetDevice(A.get_preferred_devid(k, k)));
|
|
DLAAUM(uplo, A, k, k);
|
|
}
|
|
}
|
|
|
|
void PDSYMM(cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
matrix<double>& B,
|
|
double beta,
|
|
matrix<double>& C)
|
|
{
|
|
nvtx_range r("PDSYMM");
|
|
|
|
size_t k, m, n;
|
|
double zbeta;
|
|
double zone = (double) 1.0;
|
|
|
|
for (m = 0; m < C.mt; m++)
|
|
{
|
|
for (n = 0; n < C.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(C.get_preferred_devid(m, n)));
|
|
/*
|
|
* CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER
|
|
*/
|
|
if (side == CUBLAS_SIDE_LEFT)
|
|
{
|
|
if (uplo == CUBLAS_FILL_MODE_LOWER)
|
|
{
|
|
for (k = 0; k < C.mt; k++)
|
|
{
|
|
zbeta = k == 0 ? beta : zone;
|
|
if (k < m)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, alpha, A, m, k, B, k, n, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
if (k == m)
|
|
{
|
|
DSYMM(side, uplo, alpha, A, k, k, B, k, n, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
DGEMM(CUBLAS_OP_T, CUBLAS_OP_N, alpha, A, k, m, B, k, n, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
/*
|
|
* CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_UPPER
|
|
*/
|
|
else
|
|
{
|
|
for (k = 0; k < C.mt; k++)
|
|
{
|
|
zbeta = k == 0 ? beta : zone;
|
|
if (k < m)
|
|
{
|
|
DGEMM(CUBLAS_OP_T, CUBLAS_OP_N, alpha, A, k, m, B, k, n, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
if (k == m)
|
|
{
|
|
DSYMM(side, uplo, alpha, A, k, k, B, k, n, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, alpha, A, m, k, B, k, n, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
/*
|
|
* CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_LOWER
|
|
*/
|
|
else
|
|
{
|
|
if (uplo == CUBLAS_FILL_MODE_LOWER)
|
|
{
|
|
for (k = 0; k < C.nt; k++)
|
|
{
|
|
zbeta = k == 0 ? beta : zone;
|
|
if (k < n)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, alpha, B, m, k, A, n, k, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
if (k == n)
|
|
{
|
|
DSYMM(side, uplo, alpha, A, k, k, B, m, k, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, alpha, B, m, k, A, k, n, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
/*
|
|
* CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_UPPER
|
|
*/
|
|
else
|
|
{
|
|
for (k = 0; k < C.nt; k++)
|
|
{
|
|
zbeta = k == 0 ? beta : zone;
|
|
if (k < n)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_N, alpha, B, m, k, A, k, n, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
if (k == n)
|
|
{
|
|
DSYMM(side, uplo, alpha, A, k, k, B, m, k, zbeta, C, m, n);
|
|
}
|
|
else
|
|
{
|
|
DGEMM(CUBLAS_OP_N, CUBLAS_OP_T, alpha, B, m, k, A, n, k, zbeta, C, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void PDTRMM(cublasSideMode_t side,
|
|
cublasFillMode_t uplo,
|
|
cublasOperation_t trans,
|
|
cublasDiagType_t diag,
|
|
double alpha,
|
|
matrix<double>& A,
|
|
matrix<double>& B)
|
|
{
|
|
if (side == CUBLAS_SIDE_LEFT)
|
|
{
|
|
if (uplo == CUBLAS_FILL_MODE_UPPER)
|
|
{
|
|
//===========================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_UPPER / CUBLAS_OP_N
|
|
//===========================================
|
|
if (trans == CUBLAS_OP_N)
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, m, m, B, m, n);
|
|
|
|
for (size_t k = m + 1; k < A.mt; k++)
|
|
{
|
|
DGEMM(trans, CUBLAS_OP_N, alpha, A, m, k, B, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//================================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_UPPER / CUBLAS_OP_T
|
|
//================================================
|
|
else
|
|
{
|
|
for (ssize_t m = B.mt - 1; m > -1; m--)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, m, m, B, m, n);
|
|
|
|
for (ssize_t k = 0; k < m; k++)
|
|
{
|
|
DGEMM(trans, CUBLAS_OP_N, alpha, A, k, m, B, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
//===========================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
|
|
//===========================================
|
|
if (trans == CUBLAS_OP_N)
|
|
{
|
|
for (ssize_t m = B.mt - 1; m > -1; m--)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, m, m, B, m, n);
|
|
|
|
for (ssize_t k = 0; k < m; k++)
|
|
{
|
|
DGEMM(trans, CUBLAS_OP_N, alpha, A, m, k, B, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//================================================
|
|
// CUBLAS_SIDE_LEFT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_T
|
|
//================================================
|
|
else
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
DTRMM(side, uplo, trans, diag, alpha, A, m, m, B, m, n);
|
|
|
|
for (size_t k = m + 1; k < A.mt; k++)
|
|
{
|
|
DGEMM(trans, CUBLAS_OP_N, alpha, A, k, m, B, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
if (uplo == CUBLAS_FILL_MODE_UPPER)
|
|
{
|
|
//============================================
|
|
// CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_UPPER / CUBLAS_OP_N
|
|
//============================================
|
|
if (trans == CUBLAS_OP_N)
|
|
{
|
|
for (ssize_t n = B.nt - 1; n > -1; n--)
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, n, n, B, m, n);
|
|
|
|
for (ssize_t k = 0; k < n; k++)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, trans, alpha, B, m, k, A, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//=================================================
|
|
// CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_UPPER / Plasma[_Conj]Trans
|
|
//=================================================
|
|
else
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, n, n, B, m, n);
|
|
|
|
for (size_t k = n + 1; k < A.mt; k++)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, trans, alpha, B, m, k, A, n, k, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else
|
|
{
|
|
//============================================
|
|
// CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_LOWER / CUBLAS_OP_N
|
|
//============================================
|
|
if (trans == CUBLAS_OP_N)
|
|
{
|
|
for (size_t n = 0; n < B.nt; n++)
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, n, n, B, m, n);
|
|
|
|
for (size_t k = n + 1; k < A.mt; k++)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, trans, alpha, B, m, k, A, k, n, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
//=================================================
|
|
// CUBLAS_SIDE_RIGHT / CUBLAS_FILL_MODE_LOWER / Plasma[_Conj]Trans
|
|
//=================================================
|
|
else
|
|
{
|
|
for (ssize_t n = B.nt - 1; n > -1; n--)
|
|
{
|
|
for (size_t m = 0; m < B.mt; m++)
|
|
{
|
|
cuda_try(cudaSetDevice(B.get_preferred_devid(m, n)));
|
|
|
|
DTRMM(side, uplo, trans, diag, alpha, A, n, n, B, m, n);
|
|
|
|
for (ssize_t k = 0; k < n; k++)
|
|
{
|
|
DGEMM(CUBLAS_OP_N, trans, alpha, B, m, k, A, n, k, 1.0, B, m, n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Taken from Chameleon (INRIA)
|
|
// All the formula are reported in the LAPACK Lawn 41:
|
|
// http://www.netlib.org/lapack/lawns/lawn41.ps
|
|
#define FMULS_POTRI(__n) ((double) (__n) * ((2. / 3.) + (double) (__n) * ((1. / 3.) * (double) (__n) + 1.)))
|
|
#define FADDS_POTRI(__n) ((double) (__n) * ((1. / 6.) + (double) (__n) * ((1. / 3.) * (double) (__n) - 0.5)))
|
|
double flops_dpotri(double __n)
|
|
{
|
|
double flops = (FMULS_POTRI((__n)) + FADDS_POTRI((__n)));
|
|
return flops;
|
|
}
|
|
|
|
void run(int N, int NB)
|
|
{
|
|
// Use pools of preallocated blocks
|
|
auto fixed_alloc = block_allocator<fixed_size_allocator>(ctx, NB * NB * sizeof(double));
|
|
ctx.set_allocator(fixed_alloc);
|
|
|
|
// Set up CUBLAS and CUSOLVER
|
|
int ndevs;
|
|
cuda_try(cudaGetDeviceCount(&ndevs));
|
|
|
|
for (int d = 0; d < ndevs; d++)
|
|
{
|
|
auto ldummy = ctx.logical_data(shape_of<slice<char>>(1));
|
|
ctx.task(exec_place::device(d), ldummy.write())->*[](cudaStream_t, auto) {
|
|
get_cublas_handle();
|
|
get_cusolver_handle();
|
|
};
|
|
|
|
ctx.task(exec_place::host(), ldummy.write(data_place::managed()))->*[](cudaStream_t, auto) {};
|
|
}
|
|
|
|
cuda_try(cudaSetDevice(0));
|
|
|
|
cudaStream_t timing_stream;
|
|
cuda_try(cudaStreamCreate(&timing_stream));
|
|
|
|
matrix<double> A(N, N, NB, NB, true, "A");
|
|
matrix<double> Aref(N, N, NB, NB, false, "Aref");
|
|
|
|
// (Hilbert matrix + 2*N*Id) to have a diagonal dominant matrix
|
|
auto hilbert = [=] _CCCL_HOST_DEVICE(size_t row, size_t col) {
|
|
return 1.0 / (col + row + 1.0) + 2.0 * N * (col == row);
|
|
};
|
|
|
|
Aref.fill(hilbert);
|
|
A.fill(hilbert);
|
|
|
|
/* Right-hand side */
|
|
matrix<double> B_potrs(N, 1, NB, 1, false, "B");
|
|
matrix<double> Bref_potrs(N, 1, NB, 1, false, "Bref");
|
|
|
|
auto rhs_vals = [] _CCCL_HOST_DEVICE(size_t row, size_t /*unused*/) {
|
|
return 1.0 * (row + 1);
|
|
};
|
|
|
|
B_potrs.fill(rhs_vals);
|
|
Bref_potrs.fill(rhs_vals);
|
|
|
|
int check_result = 1;
|
|
if (getenv("CHECK_RESULT"))
|
|
{
|
|
check_result = atoi(getenv("CHECK_RESULT"));
|
|
}
|
|
|
|
int check_result_potrs = check_result;
|
|
if (getenv("CHECK_RESULT_POTRS"))
|
|
{
|
|
check_result_potrs = atoi(getenv("CHECK_RESULT_POTRS"));
|
|
}
|
|
|
|
// // Compute ||Bref||
|
|
double Bref_nrm2 = 0.0;
|
|
double res_nrm2 = 0.0;
|
|
|
|
if (check_result_potrs)
|
|
{
|
|
PDNRM2_HOST(&Bref_potrs, &Bref_nrm2);
|
|
}
|
|
|
|
cudaEvent_t startEvent, stopEvent;
|
|
|
|
cuda_safe_call(cudaSetDevice(0));
|
|
cuda_safe_call(cudaStreamSynchronize(ctx.fence()));
|
|
cuda_safe_call(cudaEventCreate(&startEvent));
|
|
cuda_safe_call(cudaEventCreate(&stopEvent));
|
|
cuda_safe_call(cudaEventRecord(startEvent, ctx.fence()));
|
|
|
|
ctx.get_dot()->set_current_color("green");
|
|
PDPOTRF(A);
|
|
ctx.get_dot()->set_current_color("white");
|
|
|
|
/*
|
|
* POTRS
|
|
*/
|
|
|
|
if (check_result_potrs)
|
|
{
|
|
// Solve AX = B and put the result in B
|
|
PDPOTRS(A, B_potrs, CUBLAS_FILL_MODE_LOWER);
|
|
|
|
// Compute (AX - B)
|
|
// Bref = (Aref*B - Bref)
|
|
PDGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, Aref, B_potrs, -1.0, Bref_potrs);
|
|
|
|
// Compute ||AX - B|| = ||Bref||
|
|
PDNRM2_HOST(&Bref_potrs, &res_nrm2);
|
|
}
|
|
|
|
/*
|
|
* POTRI
|
|
*/
|
|
/* PDPOTRI = PDTRTRI + PDLAUUM */
|
|
|
|
// PDTRTRI : La^-1 (invert A)
|
|
// fprintf(stderr, "A=La before POTRI\n");
|
|
// A.print();
|
|
|
|
ctx.get_dot()->set_current_color("yellow");
|
|
PDTRTRI(A, CUBLAS_FILL_MODE_LOWER, CUBLAS_DIAG_NON_UNIT);
|
|
ctx.get_dot()->set_current_color("white");
|
|
|
|
// fprintf(stderr, "A=La^-1 after POTRI\n");
|
|
// A.print();
|
|
|
|
// Computes the lower part of A^tA (La^-t La^-1)
|
|
ctx.get_dot()->set_current_color("blue");
|
|
PDLAUUM(A, CUBLAS_FILL_MODE_LOWER);
|
|
ctx.get_dot()->set_current_color("white");
|
|
|
|
double b_nrm2_potri = 0.0;
|
|
double res_nrm2_potri = 0.0;
|
|
|
|
if (check_result)
|
|
{
|
|
/* Right-hand side */
|
|
matrix<double> B_potri(N, 1, NB, 1, false, "B_potri");
|
|
matrix<double> Bref_potri(N, 1, NB, 1, false, "Bref_potri");
|
|
|
|
// auto rhs_vals = [](matrix<double>& mat, int row, int col) { return 1.0 * (row + 1); };
|
|
B_potri.fill(rhs_vals);
|
|
Bref_potri.fill(rhs_vals);
|
|
|
|
// AX = B, X = A^-1 B
|
|
// LLt X = B, X = (LLt)^-1 B = L^-t L^-1 B
|
|
// Compute Bref_potri = (A^-1 B - B)
|
|
PDNRM2_HOST(&Bref_potri, &b_nrm2_potri);
|
|
|
|
// B = (A^-1)*B (A triangular lower, B_potri full)
|
|
// fprintf(stderr, "B_potri before PDTRMM\n");
|
|
// B_potri.print();
|
|
//
|
|
// fprintf(stderr, "A before PDTRMM\n");
|
|
// A.print();
|
|
|
|
// B_tmp = 0 (to avoid NaN*0.0)
|
|
matrix<double> B_tmp(N, 1, NB, 1, false, "B_tmp");
|
|
auto zero_vals = [] _CCCL_HOST_DEVICE(size_t /* unused */, size_t /*unused*/) {
|
|
return 0.0;
|
|
};
|
|
B_tmp.fill(zero_vals);
|
|
|
|
// B_tmp = A * B_potri + 0*B_tmp
|
|
PDSYMM(CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_LOWER, 1.0, A, B_potri, 0.0, B_tmp);
|
|
|
|
// fprintf(stderr, "B_potri after PDTRMM\n");
|
|
// B_potri.print();
|
|
|
|
// res = A X - B
|
|
PDGEMM(CUBLAS_OP_N, CUBLAS_OP_N, 1.0, Aref, B_tmp, -1.0, Bref_potri);
|
|
|
|
// fprintf(stderr, "Bref_potri after PDGEMM\n");
|
|
// Bref_potri.print();
|
|
|
|
// Compute residual
|
|
PDNRM2_HOST(&Bref_potri, &res_nrm2_potri);
|
|
}
|
|
|
|
cuda_safe_call(cudaSetDevice(0));
|
|
cuda_safe_call(cudaEventRecord(stopEvent, ctx.fence()));
|
|
|
|
ctx.finalize();
|
|
|
|
if (check_result_potrs)
|
|
{
|
|
double residual = sqrt(res_nrm2) / sqrt(Bref_nrm2);
|
|
// std::cout << "[POTRS] ||AX - B|| : " << sqrt(res_nrm2) << '\n';
|
|
// std::cout << "[POTRS] ||B|| : " << sqrt(Bref_nrm2) << '\n';
|
|
// std::cout << "[POTRS] RESIDUAL (||AX - B||/||B||) : " << residual << '\n';
|
|
EXPECT(residual < 0.01);
|
|
}
|
|
|
|
if (check_result)
|
|
{
|
|
double residual_potri = sqrt(res_nrm2_potri) / sqrt(b_nrm2_potri);
|
|
// std::cout << "[POTRI] RESIDUAL ||A * ((A^-1)B) - B|| : " << sqrt(res_nrm2_potri) << '\n';
|
|
// std::cout << "[POTRI] RESIDUAL ||B|| : " << sqrt(b_nrm2_potri) << '\n';
|
|
// std::cout << "[POTRI] RESIDUAL (||A * ((A^-1)B) - B||/||B||) : " << residual_potri << '\n';
|
|
EXPECT(residual_potri < 0.0001);
|
|
}
|
|
|
|
// // Compute Aref * A^-1 in Aref (A^-1 is lower triangular)
|
|
// PDTRMM(CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, 1.0, A, Aref);
|
|
|
|
// // This should be almost identity
|
|
// Aref.print();
|
|
|
|
#if 0
|
|
|
|
std::cout << "Print A^-1 after PDLAUUM : " << '\n';
|
|
A.print();
|
|
|
|
std::cout << "RES after AX - B POTRI : " << '\n';
|
|
Bref_potri.print();
|
|
|
|
// This should be almost identity
|
|
Aref.print();
|
|
#endif
|
|
|
|
float milliseconds;
|
|
cuda_safe_call(cudaEventElapsedTime(&milliseconds, startEvent, stopEvent));
|
|
|
|
double gflops = flops_dpotri((double) N) / (1000000000.0);
|
|
std::cout << "[PDPOTRI] ELAPSED: " << milliseconds << " ms, GFLOPS: " << gflops / (milliseconds / 1000.0) << '\n';
|
|
}
|
|
|
|
int main(int argc, char** argv)
|
|
{
|
|
int N = 1024;
|
|
int NB = 128;
|
|
|
|
if (argc > 1)
|
|
{
|
|
N = atoi(argv[1]);
|
|
}
|
|
|
|
if (argc > 2)
|
|
{
|
|
NB = atoi(argv[2]);
|
|
}
|
|
|
|
assert(N % NB == 0);
|
|
|
|
run(N, NB);
|
|
}
|