From dd7a0902160a42d402800f2ebe825e0306d7a638 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Tue, 9 Jul 2024 23:22:41 +0000 Subject: [PATCH 01/14] newer transpose-reorder --- src/tamm/CMakeLists.txt | 5 +- src/tamm/kernels/multiply.hpp | 30 +-- src/tamm/kernels/tamm_blas.hpp | 4 + src/tamm/kernels/transpose_inplace.cpp | 298 +++++++++++++++++++++++++ 4 files changed, 323 insertions(+), 14 deletions(-) create mode 100644 src/tamm/kernels/transpose_inplace.cpp diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index 31def5a65..f1c38b169 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -10,7 +10,10 @@ if(${PROJECT_NAME}_HAS_CUDA OR ${PROJECT_NAME}_HAS_HIP OR ${PROJECT_NAME}_HAS_DP mr/per_device_resource.hpp ) - set(tamm_gpu_SRCS kernels/gpu_blas.cpp) + set(tamm_gpu_SRCS + kernels/gpu_blas.cpp + kernels/transpose_inplace.cpp + ) endif() if(${PROJECT_NAME}_HAS_DPCPP) list(APPEND tamm_gpu_INCS sycl_device.hpp) diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index fc17d9c13..b6bd13527 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -15,7 +15,7 @@ #include "tamm_blas.hpp" #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) -#include "librett/librett.h" +// #include "librett/librett.h" #else namespace tamm { using gpuStream_t = int; // not used @@ -141,18 +141,22 @@ void assign_gpu(gpuStream_t& thandle, T*& dst, const SizeVec& ddims, const IntLa perm[i] = it - r_slabels.begin(); } - // create plan - librettHandle plan; -#if defined(USE_DPCPP) - sycl::queue* ptrQueue = &(thandle.first); - librettPlan(&plan, ndim, size, perm, sizeof(T), ptrQueue); -#else - librettPlan(&plan, ndim, size, perm, sizeof(T), thandle.first); -#endif - - // ABB: following casts were required since librett API only accepts void* as args - librettExecute(plan, reinterpret_cast(const_cast(src)), reinterpret_cast(dst)); - librettDestroy(plan); +// // create plan +// librettHandle plan; +// #if defined(USE_DPCPP) +// sycl::queue* ptrQueue = &(thandle.first); +// librettPlan(&plan, ndim, size, perm, sizeof(T), ptrQueue); +// #else +// librettPlan(&plan, ndim, size, perm, sizeof(T), thandle.first); +// #endif + +// // ABB: following casts were required since librett API only accepts void* as args +// librettExecute(plan, reinterpret_cast(const_cast(src)), reinterpret_cast(dst)); +// librettDestroy(plan); + + int outSize[ndim]; + for (int i = 0; i < ndim; i++) { outSize[i] = size[perm[i]]; } + gpu::transpose_inplace(dst, src, outSize, size, perm, thandle); } #endif diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index a4dca823d..053bace56 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -21,6 +21,10 @@ void axpy(const int64_t n, const T* src, const int incx, T*& dst, const int incy template void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); + +template +void transpose_inplace(T* out, const T* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); + } // namespace gpu #endif diff --git a/src/tamm/kernels/transpose_inplace.cpp b/src/tamm/kernels/transpose_inplace.cpp new file mode 100644 index 000000000..d69d2048b --- /dev/null +++ b/src/tamm/kernels/transpose_inplace.cpp @@ -0,0 +1,298 @@ +#include "tamm_blas.hpp" +#include + +namespace tamm::kernels::gpu { + +#ifdef USE_DPCPP +#define __global__ __attribute__((always_inline)) inline +#define __device__ SYCL_EXTERNAL __attribute__((always_inline)) +#endif + + template + constexpr auto divUp(U a, V b) -> decltype(a + b) { + return (a + b - 1) / b; + } + +constexpr int TILE_DIM = 32; +constexpr int THREADS_X = TILE_DIM; +constexpr int THREADS_Y = 256 / TILE_DIM; + + template + __global__ void reorderKernel(T* out, const T* in, const int* outDims, const int* inStrides, const int* outStrides, + const int d0, const int d1, const int d2, const int d3, + const int blocksPerMatX, const int blocksPerMatY +#ifdef USE_DPCPP + , sycl::nd_item<3> item +#endif + ) { + #ifndef USE_DPCPP + int threadIdx_x = threadIdx.x; + int threadIdx_y = threadIdx.y; + + int blockDim_x = blockDim.x; + int blockDim_y = blockDim.y; + + int blockIdx_x = blockIdx.x; + int blockIdx_y = blockIdx.y; + int blockIdx_z = blockIdx.z; + + int gridDim_y = gridDim.y; + #else + int threadIdx_x = item.get_local_id(2); + int threadIdx_y = item.get_local_id(1); + + int blockDim_x = item.get_local_range(2); + int blockDim_y = item.get_local_range(1); + + int blockIdx_x = item.get_group(2); + int blockIdx_y = item.get_group(1); + int blockIdx_z = item.get_group(0); + + int gridDim_y = item.get_group_range(1); + #endif + + const int oz = blockIdx_x / blocksPerMatX; + const int ow = (blockIdx_y + blockIdx_z * gridDim_y) / blocksPerMatY; + + const int blockIdx_xx = blockIdx_x - oz * blocksPerMatX; + const int blockIdx_yy = (blockIdx_y + blockIdx_z * gridDim_y) - ow * blocksPerMatY; + + const int xx = threadIdx_x + blockIdx_xx * blockDim_x; + const int yy = threadIdx_y + blockIdx_yy * blockDim_y; + + if (xx >= outDims[0] || yy >= outDims[1] || oz >= outDims[2] || + ow >= outDims[3]) + return; + + const int incy = blocksPerMatY * blockDim_y; + const int incx = blocksPerMatX * blockDim_x; + + const int rdims[] = {d0, d1, d2, d3}; + const int o_off = ow * outStrides[3] + oz * outStrides[2]; + int ids[4] = {0}; + ids[rdims[3]] = ow; + ids[rdims[2]] = oz; + + for (int oy = yy; oy < outDims[1]; oy += incy) { + ids[rdims[1]] = oy; + for (int ox = xx; ox < outDims[0]; ox += incx) { + ids[rdims[0]] = ox; + + const int oIdx = o_off + oy * outStrides[1] + ox; + + const int iIdx = ids[3] * inStrides[3] + ids[2] * inStrides[2] + + ids[1] * inStrides[1] + ids[0]; + + out[oIdx] = in[iIdx]; + } + } + } + + template + __device__ T doOp(T in, const bool conjugate) { + //return conjugate ? std::conj(in) : in; + return in; + } + + template + __global__ void transposeInPlaceKernel(T* in_ptr, + const int* in_dims, + const int* in_strides, + const int blocksPerMatX, + const int blocksPerMatY, + const bool conjugate, + const bool is32Multiple +#ifdef USE_DPCPP + , sycl::nd_item<3> it +#endif + ) { + #ifdef USE_DPCPP + sycl::group g = it.get_group(); + using tile_t = T[TILE_DIM][TILE_DIM + 1]; + tile_t& shrdMem_s = *sycl::ext::oneapi::group_local_memory_for_overwrite(g); + tile_t& shrdMem_d = *sycl::ext::oneapi::group_local_memory_for_overwrite(g); + + const int lx = it.get_local_id(2); + const int ly = it.get_local_id(1); + + // batch based block Id + const int batchId_x = g.get_group_id(2) / blocksPerMatX; + const int blockIdx_x = (g.get_group_id(2) - batchId_x * blocksPerMatX); + + const int batchId_y = g.get_group_id(1) / blocksPerMatY; + const int blockIdx_y = (g.get_group_id(1) - batchId_y * blocksPerMatY); + #else + __shared__ T shrdMem_s[TILE_DIM][TILE_DIM + 1]; + __shared__ T shrdMem_d[TILE_DIM][TILE_DIM + 1]; + + const int lx = threadIdx.x; + const int ly = threadIdx.y; + + // batch based block Id + const int batchId_x = blockIdx.x / blocksPerMatX; + const int blockIdx_x = (blockIdx.x - batchId_x * blocksPerMatX); + + const int batchId_y = blockIdx.y / blocksPerMatY; + const int blockIdx_y = (blockIdx.y - batchId_y * blocksPerMatY); + #endif + + // create variables to hold output dimensions + const int iDim0 = in_dims[0]; + const int iDim1 = in_dims[1]; + + // calculate strides + const int iStride1 = in_strides[1]; + + const int x0 = TILE_DIM * blockIdx_x; + const int y0 = TILE_DIM * blockIdx_y; + + // offset in and out based on batch id + T *iptr = in_ptr + batchId_x * in_strides[2] + batchId_y * in_strides[3]; + + if (blockIdx_y > blockIdx_x) { // Off diagonal blocks + // calculate global indices + int gx = lx + x0; + int gy = ly + y0; + int dx = lx + y0; + int dy = ly + x0; + + // Copy to shared memory + #pragma unroll + for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { + int gy_ = gy + repeat; + if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) + shrdMem_s[ly + repeat][lx] = iptr[gy_ * iStride1 + gx]; + + int dy_ = dy + repeat; + if (is32Multiple || (dx < iDim0 && dy_ < iDim1)) + shrdMem_d[ly + repeat][lx] = iptr[dy_ * iStride1 + dx]; + } + +#ifdef USE_DPCPP + sycl::group_barrier(g); +#else + __syncthreads(); +#endif + + // Copy from shared memory to global memory + #pragma unroll + for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { + int dy_ = dy + repeat; + if (is32Multiple || (dx < iDim0 && dy_ < iDim1)) + iptr[dy_ * iStride1 + dx] = doOp(shrdMem_s[lx][ly + repeat], conjugate); + + int gy_ = gy + repeat; + if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) + iptr[gy_ * iStride1 + gx] = doOp(shrdMem_d[lx][ly + repeat], conjugate); + } + + } else if (blockIdx_y == blockIdx_x) { // Diagonal blocks + // calculate global indices + int gx = lx + x0; + int gy = ly + y0; + + // offset in and out based on batch id + iptr = in_ptr + batchId_x * in_strides[2] + batchId_y * in_strides[3]; + + // Copy to shared memory + #pragma unroll + for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { + int gy_ = gy + repeat; + if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) + shrdMem_s[ly + repeat][lx] = iptr[gy_ * iStride1 + gx]; + } + +#ifdef USE_DPCPP + sycl::group_barrier(g); +#else + __syncthreads(); +#endif + + // Copy from shared memory to global memory + for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { + int gy_ = gy + repeat; + if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) + iptr[gy_ * iStride1 + gx] = doOp(shrdMem_s[lx][ly + repeat], conjugate); + } + } + } + + template + void transpose_inplace(T* out, const T* in, int* outDims, int* inDims, int* rdims, + gpuStream_t& thandle) { + + // check for conjugate here (only for complex dataypes) + bool conjugate = std::is_arithmetic_v; + + // Step 1: Launch reorder + constexpr unsigned TX = 32; + constexpr unsigned TY = 8; + constexpr unsigned TILEX = 512; + constexpr unsigned TILEY = 32; + + int blocksPerMatX = divUp(outDims[0], TILEX); + int blocksPerMatY = divUp(outDims[1], TILEY); + +#ifndef USE_DPCPP + + dim3 threads(TX, TY, 1); + dim3 blocks(blocksPerMatX * outDims[2], blocksPerMatY * outDims[3], 1); + + const int maxBlocksY = getDeviceProp(getActiveDeviceId()).maxGridSize[1]; + blocks.z = divUp(blocks.y, maxBlocksY); + blocks.y = divUp(blocks.y, blocks.z); + + reorderKernel<<>>(out, in, outDims, rdims[0], rdims[1], rdims[2], rdims[3], blocksPerMatX, blocksPerMatY); + +#else + + auto local = sycl::range{1, TY, TX}; + + sycl::id<3> max_num_wrk_groups = thandle.first.get_device().get_info>(); + auto blocks = sycl::range(1, blocksPerMatY * outDims[3], blocksPerMatX * outDims[2]); + const int maxBlocksY = max_num_wrk_groups[1]; + blocks[0] = divUp(blocks[1], maxBlocksY); + blocks[1] = divUp(blocks[1], blocks[0]); + + auto global = blocks * local; + + int stride[4] = {1,1,1,1}; + std::cout << "range of kernel launch: (" << global[0] << ", " << global[1] << ", " << global[2] << "), " << local[0] << ", " << local[1] << ", " << local[2] << std::endl; + thandle.first.parallel_for(sycl::nd_range<3>(global, local), [=](auto item) { + reorderKernel(out, in, outDims, stride, stride, rdims[0], rdims[1], rdims[2], rdims[3], blocksPerMatX, blocksPerMatY, item); + }); + +#endif + + + // Step 2: Launch transpose + const bool is32Multiple = inDims[0] % TILE_DIM == 0 && inDims[1] % TILE_DIM == 0; + + int blk_x = divUp(inDims[0], TILE_DIM); + int blk_y = divUp(inDims[1], TILE_DIM); + +#ifndef USE_DPCPP + dim3 threads(THREADS_X, THREADS_Y); + dim3 blocks(blk_x * inDims[2], blk_y * inDims[3]); + + transposeInPlaceKernel<<>>(out, outDims, in.info, blk_x, blk_y); +#else + local = sycl::range{1, THREADS_Y, THREADS_X}; + global = sycl::range(1, blk_y * THREADS_Y * inDims[3], blk_x * THREADS_X * inDims[2]); + + int outStride[4] = {1,1,1,1}; + thandle.first.parallel_for(sycl::nd_range<3>(global, local), [=](auto item) { + transposeInPlaceKernel(out, outDims, outStride, blk_x, blk_y, conjugate, is32Multiple, item); + }); +#endif + } + + + + template void transpose_inplace(double* out, const double* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); + template void transpose_inplace(float* out, const float* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); + template void transpose_inplace(std::complex* out, const std::complex* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); + template void transpose_inplace(std::complex* out, const std::complex* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); + + +} // namespace tamm::kernels::gpu From 195c222aaec14b289df690cf49fe048976abb1dd Mon Sep 17 00:00:00 2001 From: abagusetty Date: Mon, 27 Jul 2026 11:28:16 -0500 Subject: [PATCH 02/14] Replace librett with in-house reorder kernel (transpose_reorder) Replaces the librett-based tensor transpose in assign_gpu with a hand-written, rank-agnostic reorder kernel supporting CUDA/HIP/DPCPP. - New kernels/gpu_reorder.cpp: linear-index N-D axis-permuting transpose, one thread per output element. Supports ranks 1..8 (tamm::maxrank), unlike the previous 4D-only ArrayFire-derived attempt. - Column-major (first-axis-fastest) strides, verified byte-identical to librett's TensorConv convention (TensorTester reference) for all ranks, arbitrary permutations, and double/float/complex types. - Runs entirely on the caller's cuda stream / hip stream / in-order sycl queue (handle.first). No plan construction, no host round-trips, no internal synchronization -- addresses librett's synchronization issues. - Complex values transported as trivially-copyable POD bytes (device-safe; matches librett's opaque-byte handling). - Removes broken transpose_inplace.cpp (wrong in-place 2D kernel; would not compile) and the transpose_inplace declaration. - multiply.hpp: drops librett include, calls gpu::transpose_reorder. --- src/tamm/CMakeLists.txt | 2 +- src/tamm/kernels/gpu_reorder.cpp | 136 +++++++++++ src/tamm/kernels/multiply.hpp | 25 +-- src/tamm/kernels/tamm_blas.hpp | 10 +- src/tamm/kernels/transpose_inplace.cpp | 298 ------------------------- 5 files changed, 152 insertions(+), 319 deletions(-) create mode 100644 src/tamm/kernels/gpu_reorder.cpp delete mode 100644 src/tamm/kernels/transpose_inplace.cpp diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index f5f79bf28..88723225e 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -12,7 +12,7 @@ if(${PROJECT_NAME}_HAS_CUDA OR ${PROJECT_NAME}_HAS_HIP OR ${PROJECT_NAME}_HAS_DP set(tamm_gpu_SRCS kernels/gpu_blas.cpp - kernels/transpose_inplace.cpp + kernels/gpu_reorder.cpp ) endif() if(${PROJECT_NAME}_HAS_DPCPP) diff --git a/src/tamm/kernels/gpu_reorder.cpp b/src/tamm/kernels/gpu_reorder.cpp new file mode 100644 index 000000000..7b9072616 --- /dev/null +++ b/src/tamm/kernels/gpu_reorder.cpp @@ -0,0 +1,136 @@ +#include "tamm/utils.hpp" +#include "tamm_blas.hpp" + +#include + +#if defined(USE_CUDA) +#include +#elif defined(USE_HIP) +#include +#endif +// For USE_DPCPP, SYCL is pulled in transitively via tamm_blas.hpp -> +// gpu_streams.hpp -> sycl_device.hpp. + +namespace tamm::kernels::gpu { + +// Maximum tensor rank supported by the reorder kernel. Must be >= tamm::maxrank. +static constexpr int TAMM_REORDER_MAXRANK = 8; + +// POD payload passed by value into the device kernel. Holds the per-axis +// metadata needed to map an output linear index to the corresponding source +// linear index for an arbitrary axis permutation. +template +struct ReorderMeta { + int ndim; + int outDims[MAXRANK]; // extents of the (permuted) output tensor + int outStrides[MAXRANK]; // strides of the output tensor + int inStrides[MAXRANK]; // strides of the source tensor + int perm[MAXRANK]; // output-axis -> source-axis map +}; + +#if defined(USE_CUDA) || defined(USE_HIP) +// Element type used inside the device kernel. std::complex's copy assignment +// is not annotated __device__, so we transport values as a trivially-copyable +// POD of the same size/alignment (a plain byte copy of the element), exactly as +// the previous librett path did (it treated elements as opaque bytes). +template +struct DevElem { + alignas(T) unsigned char bytes[sizeof(T)]; +}; + +template +__global__ void reorder_kernel(DevElem* out, const DevElem* in, + ReorderMeta meta, size_t total) { + size_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if(tid >= total) return; + + const int ndim = meta.ndim; + + // Decode output linear index -> per-axis output index -> source linear index. + // Column-major (first-axis-fastest) layout, matching librett's convention. + size_t iIdx = 0; + for(int k = 0; k < ndim; ++k) { + int oidx = static_cast((tid / meta.outStrides[k]) % meta.outDims[k]); + // output axis k maps to source axis perm[k] + iIdx += static_cast(oidx) * meta.inStrides[meta.perm[k]]; + } + out[tid] = in[iIdx]; // trivially-copyable POD assignment (device-safe) +} +#endif + +template +void transpose_reorder(T* out, const T* in, int ndim, const int* outDims, const int* perm, + gpuStream_t& handle) { + EXPECTS(ndim >= 0 && ndim <= TAMM_REORDER_MAXRANK); + if(ndim == 0) return; + + // Source extents: inDims[perm[k]] == outDims[k]. + ReorderMeta meta{}; + meta.ndim = ndim; + + int inDims[TAMM_REORDER_MAXRANK] = {0}; + for(int k = 0; k < ndim; ++k) { + meta.outDims[k] = outDims[k]; + meta.perm[k] = perm[k]; + inDims[perm[k]] = outDims[k]; + } + + // Column-major (first-axis-fastest) contiguous strides for both tensors. + // This exactly matches librett's tensor-conversion convention (verified + // against librett's TensorTester reference), so this kernel is a drop-in + // replacement producing byte-identical output for any permutation/rank. + meta.outStrides[0] = 1; + meta.inStrides[0] = 1; + for(int k = 1; k < ndim; ++k) { + meta.outStrides[k] = meta.outStrides[k - 1] * outDims[k - 1]; + meta.inStrides[k] = meta.inStrides[k - 1] * inDims[k - 1]; + } + + size_t total = 1; + for(int k = 0; k < ndim; ++k) total *= static_cast(outDims[k]); + if(total == 0) return; + +#if defined(USE_CUDA) || defined(USE_HIP) + constexpr unsigned block = 256; + unsigned grid = static_cast((total + block - 1) / block); + reorder_kernel<<>>(reinterpret_cast*>(out), + reinterpret_cast*>(in), meta, + total); +#elif defined(USE_DPCPP) + // Transport elements as trivially-copyable POD bytes (matches the CUDA/HIP + // path and the previous librett behavior of treating elements as opaque). + struct DevElem { + alignas(T) unsigned char bytes[sizeof(T)]; + }; + auto* out_pod = reinterpret_cast(out); + const auto* in_pod = reinterpret_cast(in); + + constexpr size_t block = 256; + size_t grid = (total + block - 1) / block; + sycl::nd_range<1> ndr{sycl::range<1>(grid * block), sycl::range<1>(block)}; + handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { + size_t tid = item.get_global_id(0); + if(tid >= total) return; + + const int ndim_ = meta.ndim; + size_t iIdx = 0; + // Column-major (first-axis-fastest) layout, matching librett's convention. + for(int k = 0; k < ndim_; ++k) { + int oidx = static_cast((tid / meta.outStrides[k]) % meta.outDims[k]); + iIdx += static_cast(oidx) * meta.inStrides[meta.perm[k]]; + } + out_pod[tid] = in_pod[iIdx]; + }); +#endif +} + +template void transpose_reorder(double* out, const double* in, int ndim, const int* outDims, + const int* perm, gpuStream_t& handle); +template void transpose_reorder(float* out, const float* in, int ndim, const int* outDims, + const int* perm, gpuStream_t& handle); +template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, + const int* outDims, const int* perm, gpuStream_t& handle); +template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, + const int* outDims, const int* perm, gpuStream_t& handle); + +} // namespace tamm::kernels::gpu diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index 5f3f2d64e..2c90799bf 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -14,9 +14,7 @@ #include "tamm/utils.hpp" #include "tamm_blas.hpp" -#if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) -// #include "librett/librett.h" -#else +#if !defined(USE_CUDA) && !defined(USE_HIP) && !defined(USE_DPCPP) namespace tamm { using gpuStream_t = int; // not used } @@ -141,22 +139,11 @@ void assign_gpu(gpuStream_t& thandle, T*& dst, const SizeVec& ddims, const IntLa perm[i] = it - r_slabels.begin(); } -// // create plan -// librettHandle plan; -// #if defined(USE_DPCPP) -// sycl::queue* ptrQueue = &(thandle.first); -// librettPlan(&plan, ndim, size, perm, sizeof(T), ptrQueue); -// #else -// librettPlan(&plan, ndim, size, perm, sizeof(T), thandle.first); -// #endif - -// // ABB: following casts were required since librett API only accepts void* as args -// librettExecute(plan, reinterpret_cast(const_cast(src)), reinterpret_cast(dst)); -// librettDestroy(plan); - - int outSize[ndim]; - for (int i = 0; i < ndim; i++) { outSize[i] = size[perm[i]]; } - gpu::transpose_inplace(dst, src, outSize, size, perm, thandle); + // Out-of-place N-dimensional axis-permuting transpose (in-house reorder + // kernel; replaces the previous librett-based implementation). + int outDims[ndim]; + for(int i = 0; i < ndim; i++) { outDims[i] = size[perm[i]]; } + gpu::transpose_reorder(dst, src, ndim, outDims, perm, thandle); } #endif diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index 053bace56..39e079bbf 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -22,8 +22,16 @@ template void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); +// Out-of-place N-dimensional axis-permuting transpose (reorder). +// Replaces the previous librett-based implementation. +// out : destination buffer (device), sized as the permuted tensor +// in : source buffer (device) +// ndim : tensor rank (supports up to tamm::maxrank) +// outDims : extents of the output (permuted) tensor, length ndim +// perm : output-axis -> source-axis map, length ndim template -void transpose_inplace(T* out, const T* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); +void transpose_reorder(T* out, const T* in, int ndim, const int* outDims, const int* perm, + gpuStream_t& handle); } // namespace gpu #endif diff --git a/src/tamm/kernels/transpose_inplace.cpp b/src/tamm/kernels/transpose_inplace.cpp deleted file mode 100644 index d69d2048b..000000000 --- a/src/tamm/kernels/transpose_inplace.cpp +++ /dev/null @@ -1,298 +0,0 @@ -#include "tamm_blas.hpp" -#include - -namespace tamm::kernels::gpu { - -#ifdef USE_DPCPP -#define __global__ __attribute__((always_inline)) inline -#define __device__ SYCL_EXTERNAL __attribute__((always_inline)) -#endif - - template - constexpr auto divUp(U a, V b) -> decltype(a + b) { - return (a + b - 1) / b; - } - -constexpr int TILE_DIM = 32; -constexpr int THREADS_X = TILE_DIM; -constexpr int THREADS_Y = 256 / TILE_DIM; - - template - __global__ void reorderKernel(T* out, const T* in, const int* outDims, const int* inStrides, const int* outStrides, - const int d0, const int d1, const int d2, const int d3, - const int blocksPerMatX, const int blocksPerMatY -#ifdef USE_DPCPP - , sycl::nd_item<3> item -#endif - ) { - #ifndef USE_DPCPP - int threadIdx_x = threadIdx.x; - int threadIdx_y = threadIdx.y; - - int blockDim_x = blockDim.x; - int blockDim_y = blockDim.y; - - int blockIdx_x = blockIdx.x; - int blockIdx_y = blockIdx.y; - int blockIdx_z = blockIdx.z; - - int gridDim_y = gridDim.y; - #else - int threadIdx_x = item.get_local_id(2); - int threadIdx_y = item.get_local_id(1); - - int blockDim_x = item.get_local_range(2); - int blockDim_y = item.get_local_range(1); - - int blockIdx_x = item.get_group(2); - int blockIdx_y = item.get_group(1); - int blockIdx_z = item.get_group(0); - - int gridDim_y = item.get_group_range(1); - #endif - - const int oz = blockIdx_x / blocksPerMatX; - const int ow = (blockIdx_y + blockIdx_z * gridDim_y) / blocksPerMatY; - - const int blockIdx_xx = blockIdx_x - oz * blocksPerMatX; - const int blockIdx_yy = (blockIdx_y + blockIdx_z * gridDim_y) - ow * blocksPerMatY; - - const int xx = threadIdx_x + blockIdx_xx * blockDim_x; - const int yy = threadIdx_y + blockIdx_yy * blockDim_y; - - if (xx >= outDims[0] || yy >= outDims[1] || oz >= outDims[2] || - ow >= outDims[3]) - return; - - const int incy = blocksPerMatY * blockDim_y; - const int incx = blocksPerMatX * blockDim_x; - - const int rdims[] = {d0, d1, d2, d3}; - const int o_off = ow * outStrides[3] + oz * outStrides[2]; - int ids[4] = {0}; - ids[rdims[3]] = ow; - ids[rdims[2]] = oz; - - for (int oy = yy; oy < outDims[1]; oy += incy) { - ids[rdims[1]] = oy; - for (int ox = xx; ox < outDims[0]; ox += incx) { - ids[rdims[0]] = ox; - - const int oIdx = o_off + oy * outStrides[1] + ox; - - const int iIdx = ids[3] * inStrides[3] + ids[2] * inStrides[2] + - ids[1] * inStrides[1] + ids[0]; - - out[oIdx] = in[iIdx]; - } - } - } - - template - __device__ T doOp(T in, const bool conjugate) { - //return conjugate ? std::conj(in) : in; - return in; - } - - template - __global__ void transposeInPlaceKernel(T* in_ptr, - const int* in_dims, - const int* in_strides, - const int blocksPerMatX, - const int blocksPerMatY, - const bool conjugate, - const bool is32Multiple -#ifdef USE_DPCPP - , sycl::nd_item<3> it -#endif - ) { - #ifdef USE_DPCPP - sycl::group g = it.get_group(); - using tile_t = T[TILE_DIM][TILE_DIM + 1]; - tile_t& shrdMem_s = *sycl::ext::oneapi::group_local_memory_for_overwrite(g); - tile_t& shrdMem_d = *sycl::ext::oneapi::group_local_memory_for_overwrite(g); - - const int lx = it.get_local_id(2); - const int ly = it.get_local_id(1); - - // batch based block Id - const int batchId_x = g.get_group_id(2) / blocksPerMatX; - const int blockIdx_x = (g.get_group_id(2) - batchId_x * blocksPerMatX); - - const int batchId_y = g.get_group_id(1) / blocksPerMatY; - const int blockIdx_y = (g.get_group_id(1) - batchId_y * blocksPerMatY); - #else - __shared__ T shrdMem_s[TILE_DIM][TILE_DIM + 1]; - __shared__ T shrdMem_d[TILE_DIM][TILE_DIM + 1]; - - const int lx = threadIdx.x; - const int ly = threadIdx.y; - - // batch based block Id - const int batchId_x = blockIdx.x / blocksPerMatX; - const int blockIdx_x = (blockIdx.x - batchId_x * blocksPerMatX); - - const int batchId_y = blockIdx.y / blocksPerMatY; - const int blockIdx_y = (blockIdx.y - batchId_y * blocksPerMatY); - #endif - - // create variables to hold output dimensions - const int iDim0 = in_dims[0]; - const int iDim1 = in_dims[1]; - - // calculate strides - const int iStride1 = in_strides[1]; - - const int x0 = TILE_DIM * blockIdx_x; - const int y0 = TILE_DIM * blockIdx_y; - - // offset in and out based on batch id - T *iptr = in_ptr + batchId_x * in_strides[2] + batchId_y * in_strides[3]; - - if (blockIdx_y > blockIdx_x) { // Off diagonal blocks - // calculate global indices - int gx = lx + x0; - int gy = ly + y0; - int dx = lx + y0; - int dy = ly + x0; - - // Copy to shared memory - #pragma unroll - for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { - int gy_ = gy + repeat; - if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) - shrdMem_s[ly + repeat][lx] = iptr[gy_ * iStride1 + gx]; - - int dy_ = dy + repeat; - if (is32Multiple || (dx < iDim0 && dy_ < iDim1)) - shrdMem_d[ly + repeat][lx] = iptr[dy_ * iStride1 + dx]; - } - -#ifdef USE_DPCPP - sycl::group_barrier(g); -#else - __syncthreads(); -#endif - - // Copy from shared memory to global memory - #pragma unroll - for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { - int dy_ = dy + repeat; - if (is32Multiple || (dx < iDim0 && dy_ < iDim1)) - iptr[dy_ * iStride1 + dx] = doOp(shrdMem_s[lx][ly + repeat], conjugate); - - int gy_ = gy + repeat; - if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) - iptr[gy_ * iStride1 + gx] = doOp(shrdMem_d[lx][ly + repeat], conjugate); - } - - } else if (blockIdx_y == blockIdx_x) { // Diagonal blocks - // calculate global indices - int gx = lx + x0; - int gy = ly + y0; - - // offset in and out based on batch id - iptr = in_ptr + batchId_x * in_strides[2] + batchId_y * in_strides[3]; - - // Copy to shared memory - #pragma unroll - for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { - int gy_ = gy + repeat; - if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) - shrdMem_s[ly + repeat][lx] = iptr[gy_ * iStride1 + gx]; - } - -#ifdef USE_DPCPP - sycl::group_barrier(g); -#else - __syncthreads(); -#endif - - // Copy from shared memory to global memory - for (int repeat = 0; repeat < TILE_DIM; repeat += THREADS_Y) { - int gy_ = gy + repeat; - if (is32Multiple || (gx < iDim0 && gy_ < iDim1)) - iptr[gy_ * iStride1 + gx] = doOp(shrdMem_s[lx][ly + repeat], conjugate); - } - } - } - - template - void transpose_inplace(T* out, const T* in, int* outDims, int* inDims, int* rdims, - gpuStream_t& thandle) { - - // check for conjugate here (only for complex dataypes) - bool conjugate = std::is_arithmetic_v; - - // Step 1: Launch reorder - constexpr unsigned TX = 32; - constexpr unsigned TY = 8; - constexpr unsigned TILEX = 512; - constexpr unsigned TILEY = 32; - - int blocksPerMatX = divUp(outDims[0], TILEX); - int blocksPerMatY = divUp(outDims[1], TILEY); - -#ifndef USE_DPCPP - - dim3 threads(TX, TY, 1); - dim3 blocks(blocksPerMatX * outDims[2], blocksPerMatY * outDims[3], 1); - - const int maxBlocksY = getDeviceProp(getActiveDeviceId()).maxGridSize[1]; - blocks.z = divUp(blocks.y, maxBlocksY); - blocks.y = divUp(blocks.y, blocks.z); - - reorderKernel<<>>(out, in, outDims, rdims[0], rdims[1], rdims[2], rdims[3], blocksPerMatX, blocksPerMatY); - -#else - - auto local = sycl::range{1, TY, TX}; - - sycl::id<3> max_num_wrk_groups = thandle.first.get_device().get_info>(); - auto blocks = sycl::range(1, blocksPerMatY * outDims[3], blocksPerMatX * outDims[2]); - const int maxBlocksY = max_num_wrk_groups[1]; - blocks[0] = divUp(blocks[1], maxBlocksY); - blocks[1] = divUp(blocks[1], blocks[0]); - - auto global = blocks * local; - - int stride[4] = {1,1,1,1}; - std::cout << "range of kernel launch: (" << global[0] << ", " << global[1] << ", " << global[2] << "), " << local[0] << ", " << local[1] << ", " << local[2] << std::endl; - thandle.first.parallel_for(sycl::nd_range<3>(global, local), [=](auto item) { - reorderKernel(out, in, outDims, stride, stride, rdims[0], rdims[1], rdims[2], rdims[3], blocksPerMatX, blocksPerMatY, item); - }); - -#endif - - - // Step 2: Launch transpose - const bool is32Multiple = inDims[0] % TILE_DIM == 0 && inDims[1] % TILE_DIM == 0; - - int blk_x = divUp(inDims[0], TILE_DIM); - int blk_y = divUp(inDims[1], TILE_DIM); - -#ifndef USE_DPCPP - dim3 threads(THREADS_X, THREADS_Y); - dim3 blocks(blk_x * inDims[2], blk_y * inDims[3]); - - transposeInPlaceKernel<<>>(out, outDims, in.info, blk_x, blk_y); -#else - local = sycl::range{1, THREADS_Y, THREADS_X}; - global = sycl::range(1, blk_y * THREADS_Y * inDims[3], blk_x * THREADS_X * inDims[2]); - - int outStride[4] = {1,1,1,1}; - thandle.first.parallel_for(sycl::nd_range<3>(global, local), [=](auto item) { - transposeInPlaceKernel(out, outDims, outStride, blk_x, blk_y, conjugate, is32Multiple, item); - }); -#endif - } - - - - template void transpose_inplace(double* out, const double* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); - template void transpose_inplace(float* out, const float* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); - template void transpose_inplace(std::complex* out, const std::complex* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); - template void transpose_inplace(std::complex* out, const std::complex* in, int* outDims, int* inDims, int* rdims, gpuStream_t& thandle); - - -} // namespace tamm::kernels::gpu From 6b35ec0e115a108276b0a1e1aae28eff3b2fea26 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Mon, 21 Sep 2026 11:03:31 -0500 Subject: [PATCH 03/14] Replace librett with in-house GPU transpose, fix bufacc overwrite Drop the librett dependency: new transpose_reorder kernel (CUDA/HIP/SYCL) driven by allocation-free ReorderSpec metadata. GPU output transpose keeps librett overwrite semantics (GEMM beta accumulates the running total, so accumulating again in the transpose double-counts). Full ExaChem CI 19/19. --- CMakeLists.txt | 8 - docs/install.rst | 2 +- src/tamm/CMakeLists.txt | 8 + src/tamm/kernels/gpu_blas.cpp | 13 +- src/tamm/kernels/gpu_reorder.cpp | 253 ++++++++++++++++----------- src/tamm/kernels/gpu_reorder.hpp | 225 ++++++++++++++++++++++++ src/tamm/kernels/multiply.hpp | 88 +++++----- src/tamm/kernels/tamm_blas.hpp | 24 ++- tests/tamm/Test_TransposeReorder.cpp | 239 +++++++++++++++++++++++++ tests/tamm/test_tamm.cmake | 1 + 10 files changed, 689 insertions(+), 172 deletions(-) create mode 100644 src/tamm/kernels/gpu_reorder.hpp create mode 100644 tests/tamm/Test_TransposeReorder.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index e0acd6030..93265cad2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,14 +50,6 @@ elseif(NOT DEFINED USE_LIBNUMA OR USE_LIBNUMA) list(APPEND TAMM_DEPENDENCIES numactl) endif() -if(TAMM_ENABLE_CUDA OR TAMM_ENABLE_HIP OR TAMM_ENABLE_DPCPP) - if(USE_SYCL_BLAS) - #list(APPEND TAMM_DEPENDENCIES sycl-blas) - endif() - - list(APPEND TAMM_DEPENDENCIES Librett) -endif() - if(TAMM_ENABLE_ELPA) list(APPEND TAMM_DEPENDENCIES ELPA) endif() diff --git a/docs/install.rst b/docs/install.rst index 384841c4e..c200430ed 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -17,7 +17,7 @@ Dependencies **The remaining dependencies are automatically built and do not need to be installed explicitly:** * GlobalArrays -* HPTT, Librett +* HPTT * HDF5 * BLAS/LAPACK (BLIS and netlib-lapack are automatically built if no vendor BLAS libraries are provided) * BLAS++ and LAPACK++ diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index 09a614519..af76e2e78 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -15,6 +15,13 @@ if(${PROJECT_NAME}_HAS_CUDA OR ${PROJECT_NAME}_HAS_HIP OR ${PROJECT_NAME}_HAS_DP kernels/gpu_lapack.cpp kernels/gpu_reorder.cpp ) + if(${PROJECT_NAME}_HAS_CUDA) + # gpu_reorder.cpp holds __global__ kernels + <<<>>> launches, so nvcc + # (not the C++ compiler) must compile it even though it is a .cpp file. + set_source_files_properties(kernels/gpu_reorder.cpp PROPERTIES LANGUAGE CUDA) + endif() + # NOTE: hipcc compiles .cpp HIP code directly, and icpx compiles the SYCL + # branch, so no LANGUAGE override is needed for HIP/DPCPP builds. endif() if(${PROJECT_NAME}_HAS_DPCPP) list(APPEND tamm_gpu_INCS sycl_device.hpp) @@ -77,6 +84,7 @@ set(TAMM_INCLUDES kernels/assign.hpp kernels/multiply.hpp kernels/tamm_blas.hpp + kernels/gpu_reorder.hpp op_dag.hpp interfaces.hpp op_profiler.hpp diff --git a/src/tamm/kernels/gpu_blas.cpp b/src/tamm/kernels/gpu_blas.cpp index 1ec16be38..8b7513f92 100644 --- a/src/tamm/kernels/gpu_blas.cpp +++ b/src/tamm/kernels/gpu_blas.cpp @@ -42,19 +42,20 @@ void tamm::kernels::gpu::axpy(const int64_t n, const T* src, const int incx, T*& template void tamm::kernels::gpu::gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, - int lda, const T beta, T1* C, int ldc, gpuStream_t& handle) { + int lda, const T beta, T1* C, int ldc, gpuStream_t& handle) { #if defined(USE_DPCPP) #ifdef USE_PORT_BLAS blas::SB_Handle sb_handle(handle.first); blas::internal::_gemm(sb_handle, 'n', 'n', n, m, k, alpha, const_cast(B), ldb, const_cast(A), lda, beta, C, ldc, {}); - handle.first.wait(); + // No wait: the caller's in-order queue preserves ordering and + // block_multiply() synchronizes before pool buffers are recycled. #else - auto gemm_event = oneapi::mkl::blas::column_major::gemm(handle.first, oneapi::mkl::transpose::N, - oneapi::mkl::transpose::N, n, m, k, alpha, - B, ldb, A, lda, beta, C, ldc); - gemm_event.wait(); + // No wait (same reason); the event is intentionally discarded. + oneapi::mkl::blas::column_major::gemm(handle.first, oneapi::mkl::transpose::N, + oneapi::mkl::transpose::N, n, m, k, alpha, B, ldb, A, + lda, beta, C, ldc); #endif // USE_PORT_BLAS #elif defined(USE_CUDA) diff --git a/src/tamm/kernels/gpu_reorder.cpp b/src/tamm/kernels/gpu_reorder.cpp index 7b9072616..8ba766c9f 100644 --- a/src/tamm/kernels/gpu_reorder.cpp +++ b/src/tamm/kernels/gpu_reorder.cpp @@ -1,7 +1,10 @@ -#include "tamm/utils.hpp" +#include "gpu_reorder.hpp" + #include "tamm_blas.hpp" #include +#include +#include #if defined(USE_CUDA) #include @@ -13,124 +16,174 @@ namespace tamm::kernels::gpu { -// Maximum tensor rank supported by the reorder kernel. Must be >= tamm::maxrank. -static constexpr int TAMM_REORDER_MAXRANK = 8; - -// POD payload passed by value into the device kernel. Holds the per-axis -// metadata needed to map an output linear index to the corresponding source -// linear index for an arbitrary axis permutation. -template -struct ReorderMeta { - int ndim; - int outDims[MAXRANK]; // extents of the (permuted) output tensor - int outStrides[MAXRANK]; // strides of the output tensor - int inStrides[MAXRANK]; // strides of the source tensor - int perm[MAXRANK]; // output-axis -> source-axis map -}; - -#if defined(USE_CUDA) || defined(USE_HIP) -// Element type used inside the device kernel. std::complex's copy assignment -// is not annotated __device__, so we transport values as a trivially-copyable -// POD of the same size/alignment (a plain byte copy of the element), exactly as -// the previous librett path did (it treated elements as opaque bytes). +// Decompose the host-side scale into raw {real, imag} doubles for the kernel. template -struct DevElem { - alignas(T) unsigned char bytes[sizeof(T)]; -}; +inline void reorder_split_scale(T scale, double& re, double& im) { + if constexpr(reorder_is_complex_v) { + re = static_cast(scale.real()); + im = static_cast(scale.imag()); + } + else { + re = static_cast(scale); + im = 0.0; + } +} template -__global__ void reorder_kernel(DevElem* out, const DevElem* in, - ReorderMeta meta, size_t total) { - size_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if(tid >= total) return; - - const int ndim = meta.ndim; - - // Decode output linear index -> per-axis output index -> source linear index. - // Column-major (first-axis-fastest) layout, matching librett's convention. - size_t iIdx = 0; - for(int k = 0; k < ndim; ++k) { - int oidx = static_cast((tid / meta.outStrides[k]) % meta.outDims[k]); - // output axis k maps to source axis perm[k] - iIdx += static_cast(oidx) * meta.inStrides[meta.perm[k]]; +inline bool reorder_scale_is_one(T scale) { + if constexpr(reorder_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } + else { return scale == T{1}; } +} + +#if defined(USE_CUDA) || defined(USE_HIP) +// One thread per output element, grid-stride loop: any total works with a +// bounded grid, and the kernel never touches the caller's stream ordering +// (no internal sync, no plan, no host round-trip). +template +__global__ void +#if defined(USE_CUDA) || defined(USE_HIP) +__launch_bounds__(256) +#endif +reorder_kernel(T* out, const T* in, ReorderMeta meta, size_t total, double scale_re, double scale_im, + bool accumulate) +{ + const size_t tid0 = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t stride = static_cast(gridDim.x) * blockDim.x; + for(size_t t = tid0; t < total; t += stride) { + const Idx tid = static_cast(t); + const size_t src = reorder_src_index(tid, meta); + const T y = reorder_scaled(in[src], scale_re, scale_im); + out[tid] = accumulate ? reorder_add(out[tid], y) : y; } - out[tid] = in[iIdx]; // trivially-copyable POD assignment (device-safe) } + +inline void reorder_check_launch() { +#if defined(USE_CUDA) + const cudaError_t err = cudaGetLastError(); + if(err != cudaSuccess) { + std::ostringstream msg; + msg << "TAMM reorder kernel launch failed: " << cudaGetErrorString(err); + throw std::runtime_error(msg.str()); + } +#elif defined(USE_HIP) + const hipError_t err = hipGetLastError(); + if(err != hipSuccess) { + std::ostringstream msg; + msg << "TAMM reorder kernel launch failed: " << hipGetErrorString(err); + throw std::runtime_error(msg.str()); + } #endif +} +#endif // USE_CUDA || USE_HIP template -void transpose_reorder(T* out, const T* in, int ndim, const int* outDims, const int* perm, - gpuStream_t& handle) { - EXPECTS(ndim >= 0 && ndim <= TAMM_REORDER_MAXRANK); - if(ndim == 0) return; - - // Source extents: inDims[perm[k]] == outDims[k]. - ReorderMeta meta{}; - meta.ndim = ndim; - - int inDims[TAMM_REORDER_MAXRANK] = {0}; - for(int k = 0; k < ndim; ++k) { - meta.outDims[k] = outDims[k]; - meta.perm[k] = perm[k]; - inDims[perm[k]] = outDims[k]; +void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, + T scale, bool accumulate, gpuStream_t& handle) { + EXPECTS(ndim >= 0 && ndim <= reorder_maxrank); + EXPECTS(out != in); // out-of-place only; in-place permutes need a temp buffer + if(ndim == 0) { + // Rank-0: single element, still honors scale/accumulate. + double scale_re, scale_im; + reorder_split_scale(scale, scale_re, scale_im); +#if defined(USE_CUDA) || defined(USE_HIP) + // Route through the kernel for stream ordering (async, like every + // other path through this function). + if(reorder_scale_is_one(scale) && !accumulate) { + gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); + } + else { + ReorderMeta meta{}; + meta.ndim = 0; + reorder_kernel + <<<1, 1, 0, handle.first>>>(out, in, meta, 1, scale_re, scale_im, accumulate); + reorder_check_launch(); + } +#elif defined(USE_DPCPP) + if(reorder_scale_is_one(scale) && !accumulate) { + gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); + } + else { + handle.first.parallel_for(sycl::range<1>(1), [=](sycl::id<1>) { + const T y = reorder_scaled(in[0], scale_re, scale_im); + out[0] = accumulate ? reorder_add(out[0], y) : y; + }); + } +#endif + return; } + EXPECTS(outDims != nullptr && perm != nullptr); + + const ReorderMeta meta = reorder_build_meta(ndim, outDims, perm); + const size_t total = reorder_total(ndim, outDims); + if(total == 0) return; + + double scale_re, scale_im; + reorder_split_scale(scale, scale_re, scale_im); + if constexpr(!reorder_is_complex_v) { EXPECTS(scale_im == 0.0); } - // Column-major (first-axis-fastest) contiguous strides for both tensors. - // This exactly matches librett's tensor-conversion convention (verified - // against librett's TensorTester reference), so this kernel is a drop-in - // replacement producing byte-identical output for any permutation/rank. - meta.outStrides[0] = 1; - meta.inStrides[0] = 1; - for(int k = 1; k < ndim; ++k) { - meta.outStrides[k] = meta.outStrides[k - 1] * outDims[k - 1]; - meta.inStrides[k] = meta.inStrides[k - 1] * inDims[k - 1]; + // Identity permutation with scale==1 and overwrite: pure copy, no kernel. + if(!accumulate && reorder_scale_is_one(scale) && reorder_is_identity(meta)) { + gpuMemcpyAsync(out, in, total, gpuMemcpyDeviceToDevice, handle); + return; } - size_t total = 1; - for(int k = 0; k < ndim; ++k) total *= static_cast(outDims[k]); - if(total == 0) return; + constexpr size_t block = 256; + constexpr size_t maxBlocks = 1 << 20; // stride loop covers the rest + const size_t nblocks = std::min((total + block - 1) / block, maxBlocks); #if defined(USE_CUDA) || defined(USE_HIP) - constexpr unsigned block = 256; - unsigned grid = static_cast((total + block - 1) / block); - reorder_kernel<<>>(reinterpret_cast*>(out), - reinterpret_cast*>(in), meta, - total); + if(reorder_meta_fits32(meta, total)) { + reorder_kernel<<(nblocks), static_cast(block), 0, + handle.first>>>(out, in, meta, total, scale_re, scale_im, + accumulate); + } + else { + reorder_kernel<<(nblocks), static_cast(block), 0, + handle.first>>>(out, in, meta, total, scale_re, scale_im, + accumulate); + } + reorder_check_launch(); #elif defined(USE_DPCPP) - // Transport elements as trivially-copyable POD bytes (matches the CUDA/HIP - // path and the previous librett behavior of treating elements as opaque). - struct DevElem { - alignas(T) unsigned char bytes[sizeof(T)]; - }; - auto* out_pod = reinterpret_cast(out); - const auto* in_pod = reinterpret_cast(in); - - constexpr size_t block = 256; - size_t grid = (total + block - 1) / block; - sycl::nd_range<1> ndr{sycl::range<1>(grid * block), sycl::range<1>(block)}; - handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { - size_t tid = item.get_global_id(0); - if(tid >= total) return; - - const int ndim_ = meta.ndim; - size_t iIdx = 0; - // Column-major (first-axis-fastest) layout, matching librett's convention. - for(int k = 0; k < ndim_; ++k) { - int oidx = static_cast((tid / meta.outStrides[k]) % meta.outDims[k]); - iIdx += static_cast(oidx) * meta.inStrides[meta.perm[k]]; - } - out_pod[tid] = in_pod[iIdx]; - }); + const sycl::nd_range<1> ndr(sycl::range<1>(nblocks * block), sycl::range<1>(block)); + if(reorder_meta_fits32(meta, total)) { + handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { + const size_t tid0 = item.get_global_id(0); + const size_t stride = item.get_global_range(0); + for(size_t t = tid0; t < total; t += stride) { + const uint32_t tid = static_cast(t); + const size_t src = reorder_src_index(tid, meta); + const T y = reorder_scaled(in[src], scale_re, scale_im); + out[tid] = accumulate ? reorder_add(out[tid], y) : y; + } + }); + } + else { + handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { + const size_t tid0 = item.get_global_id(0); + const size_t stride = item.get_global_range(0); + for(size_t t = tid0; t < total; t += stride) { + const size_t src = reorder_src_index(t, meta); + const T y = reorder_scaled(in[src], scale_re, scale_im); + out[t] = accumulate ? reorder_add(out[t], y) : y; + } + }); + } + // No wait: the caller's in-order queue preserves ordering, and + // block_multiply() synchronizes before pool buffers are recycled. #endif } -template void transpose_reorder(double* out, const double* in, int ndim, const int* outDims, - const int* perm, gpuStream_t& handle); -template void transpose_reorder(float* out, const float* in, int ndim, const int* outDims, - const int* perm, gpuStream_t& handle); +template void transpose_reorder(double* out, const double* in, int ndim, const size_t* outDims, + const int* perm, double scale, bool accumulate, + gpuStream_t& handle); +template void transpose_reorder(float* out, const float* in, int ndim, const size_t* outDims, + const int* perm, float scale, bool accumulate, + gpuStream_t& handle); template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, - const int* outDims, const int* perm, gpuStream_t& handle); + const size_t* outDims, const int* perm, std::complex scale, + bool accumulate, gpuStream_t& handle); template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, - const int* outDims, const int* perm, gpuStream_t& handle); + const size_t* outDims, const int* perm, std::complex scale, + bool accumulate, gpuStream_t& handle); } // namespace tamm::kernels::gpu diff --git a/src/tamm/kernels/gpu_reorder.hpp b/src/tamm/kernels/gpu_reorder.hpp new file mode 100644 index 000000000..b9148e881 --- /dev/null +++ b/src/tamm/kernels/gpu_reorder.hpp @@ -0,0 +1,225 @@ +#pragma once + +// Host/device shared logic for TAMM's in-house GPU tensor transpose (reorder). +// +// This header holds everything about the transpose that is backend-agnostic: +// - the permutation metadata layout, +// - the output-linear-index -> input-linear-index decode (column-major, +// first-axis-fastest, matching the convention the old librett path used), +// - the scale/accumulate elementwise apply, +// - host-side helpers that build the metadata from TAMM's (dims, labels). +// +// The per-backend kernels in gpu_reorder.cpp are thin launch wrappers around +// these helpers, so all three backends (CUDA/HIP/SYCL) provably implement the +// same math. The CPU unit test (tests/tamm/Test_TransposeReorder.cpp) +// exercises these helpers directly, without needing a GPU. +// +// Layout note: TAMM callers pass dims/labels in reverse (Fortran) order +// (see build_reorder_spec); the kernel itself always works in column-major +// (first-axis-fastest) order. Reverse + column-major is exactly equivalent to +// the natural row-major (last-axis-fastest) transpose, which is what the CPU +// reference (kernels::assign / HPTT path) computes. + +#include +#include +#include +#include +#include +#include +#include + +#include "tamm/errors.hpp" +#include "tamm/types.hpp" + +#if defined(__CUDACC__) || defined(__HIPCC__) +#define TAMM_REORDER_HD __host__ __device__ +#else +#define TAMM_REORDER_HD +#endif + +namespace tamm::kernels::gpu { + +// Maximum tensor rank supported by the reorder kernel. +inline constexpr int reorder_maxrank = 8; +static_assert(reorder_maxrank >= static_cast(tamm::maxrank), + "reorder kernel must support at least tamm::maxrank"); + +template +struct reorder_is_complex: std::false_type {}; +template +struct reorder_is_complex>: std::true_type {}; +template +inline constexpr bool reorder_is_complex_v = reorder_is_complex::value; + +// POD metadata passed by value into the device kernel. Maps an output linear +// index to the corresponding source linear index for an arbitrary axis +// permutation. All extents/strides are 64-bit: TAMM blocks can exceed 2^31 +// elements (cf. the B*K*N = 2^32 overflow previously fixed in gemm_wrapper). +struct ReorderMeta { + int ndim = 0; + size_t outDims[reorder_maxrank] = {}; + size_t outStrides[reorder_maxrank] = {}; + size_t inStrides[reorder_maxrank] = {}; + int perm[reorder_maxrank] = {}; // output-axis -> source-axis map +}; +// Passed by value into CUDA/HIP kernels (4 KB parameter limit) and captured +// by value into SYCL kernels (must stay trivially copyable). +static_assert(sizeof(ReorderMeta) <= 1024, "ReorderMeta must stay small"); +static_assert(std::is_trivially_copyable_v, + "ReorderMeta must be trivially copyable for device kernels"); + +// Decode: output linear index -> source linear index (column-major). +// Templated on the integer type so launches with < 2^32 elements can use +// 32-bit division/modulo (much faster on GPUs); the caller guarantees every +// value used below fits in Idx (see meta_fits32). +template +TAMM_REORDER_HD inline size_t reorder_src_index(Idx tid, const ReorderMeta& meta) { + size_t src = 0; + for(int k = 0; k < meta.ndim; ++k) { + const Idx oidx = + static_cast((tid / static_cast(meta.outStrides[k])) % + static_cast(meta.outDims[k])); + src += static_cast(oidx) * meta.inStrides[meta.perm[k]]; + } + return src; +} + +// Elementwise math, split in two so device code never default-constructs a +// T (std::complex's default ctor is not reliably usable in CUDA/HIP/SYCL +// device code) and the overwrite path never reads the output buffer. +// Complex arithmetic runs on the raw {real, imag} words so no std::complex +// device calls are needed; the complex <-> T[2] pun is blessed by the +// standard. +template +TAMM_REORDER_HD inline T reorder_scaled(T v, double scale_re, double scale_im) { + if constexpr(reorder_is_complex_v) { + using R = typename T::value_type; + const R* w = reinterpret_cast(&v); + T out; + R* o = reinterpret_cast(&out); + o[0] = static_cast(scale_re * w[0] - scale_im * w[1]); + o[1] = static_cast(scale_re * w[1] + scale_im * w[0]); + return out; + } + else { + return static_cast(scale_re * static_cast(v)); + } +} + +template +TAMM_REORDER_HD inline T reorder_add(T a, T b) { + if constexpr(reorder_is_complex_v) { + using R = typename T::value_type; + const R* wa = reinterpret_cast(&a); + const R* wb = reinterpret_cast(&b); + T out; + R* o = reinterpret_cast(&out); + o[0] = wa[0] + wb[0]; + o[1] = wa[1] + wb[1]; + return out; + } + else { + return a + b; + } +} + +// Per-element store idiom (use at every call site): +// const T y = reorder_scaled(in[src], scale_re, scale_im); +// out[tid] = accumulate ? reorder_add(out[tid], y) : y; +// so the overwrite path never reads the output buffer and no code path +// default-constructs a T on the device. + +// True when every extent/stride/total used by the decode fits in 32 bits. +inline bool reorder_meta_fits32(const ReorderMeta& meta, size_t total) { + constexpr auto max32 = static_cast(std::numeric_limits::max()); + if(total > max32) return false; + for(int k = 0; k < meta.ndim; ++k) { + if(meta.outDims[k] > max32 || meta.outStrides[k] > max32 || meta.inStrides[k] > max32) { + return false; + } + } + return true; +} + +inline bool reorder_is_identity(const ReorderMeta& meta) { + for(int k = 0; k < meta.ndim; ++k) { + if(meta.perm[k] != k) return false; + } + return true; +} + +// Host-side: build the metadata from the (already reversed) output extents +// and output-axis -> source-axis permutation. Source extents follow from +// inDims[perm[k]] == outDims[k]; both tensors are contiguous column-major. +inline ReorderMeta reorder_build_meta(int ndim, const size_t* outDims, const int* perm) { + EXPECTS(ndim >= 0 && ndim <= reorder_maxrank); + ReorderMeta meta{}; + meta.ndim = ndim; + if(ndim == 0) return meta; + + size_t inDims[reorder_maxrank] = {}; + for(int k = 0; k < ndim; ++k) { + EXPECTS(perm[k] >= 0 && perm[k] < ndim); + meta.outDims[k] = outDims[k]; + meta.perm[k] = perm[k]; + inDims[perm[k]] = outDims[k]; + } + meta.outStrides[0] = 1; + meta.inStrides[0] = 1; + for(int k = 1; k < ndim; ++k) { + meta.outStrides[k] = meta.outStrides[k - 1] * outDims[k - 1]; + meta.inStrides[k] = meta.inStrides[k - 1] * inDims[k - 1]; + } + return meta; +} + +inline size_t reorder_total(int ndim, const size_t* outDims) { + size_t total = 1; + for(int k = 0; k < ndim; ++k) total *= outDims[k]; + return total; +} + +// Host-side: translate TAMM's (src dims, src labels, dst labels) into the +// reversed (Fortran-order) output extents + permutation the column-major +// kernel consumes. Reversing both dim orders converts between TAMM's natural +// row-major transpose and the kernel's column-major transpose. +// +// Allocation-free: everything lives in fixed-size stack arrays (rank <= +// reorder_maxrank), so this can sit in the hottest transpose path without +// touching the heap. (An earlier version used std::vectors here: 4 heap +// allocations per transpose call.) +struct ReorderSpec { + int ndim = 0; + size_t outDims[reorder_maxrank] = {}; + int perm[reorder_maxrank] = {}; // reversed output-axis -> source-axis map +}; + +inline void build_reorder_spec(const SizeVec& sdims, const IntLabelVec& slabels, + const IntLabelVec& dlabels, ReorderSpec& spec) { + const size_t ndim = sdims.size(); + EXPECTS(dlabels.size() == ndim && slabels.size() == ndim); + EXPECTS(ndim <= static_cast(reorder_maxrank)); + + // Reverse into Fortran order (librett used the same convention). + size_t r_size[reorder_maxrank]; + IntLabel r_slabels[reorder_maxrank]; + IntLabel r_dlabels[reorder_maxrank]; + for(size_t i = 0; i < ndim; ++i) { + r_size[i] = sdims[ndim - 1 - i].value(); + r_slabels[i] = slabels[ndim - 1 - i]; + r_dlabels[i] = dlabels[ndim - 1 - i]; + } + + spec.ndim = static_cast(ndim); + for(size_t i = 0; i < ndim; ++i) { + int j = 0; + while(j < static_cast(ndim) && r_slabels[static_cast(j)] != r_dlabels[i]) ++j; + EXPECTS(j < static_cast(ndim)); + spec.perm[i] = j; + } + for(size_t i = 0; i < ndim; ++i) { + spec.outDims[i] = r_size[static_cast(spec.perm[i])]; + } +} + +} // namespace tamm::kernels::gpu diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index 42e9836d5..3220c59c3 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -15,6 +15,7 @@ #include "tamm/rmm_memory_pool.hpp" #include "tamm/utils.hpp" #include "tamm_blas.hpp" +#include "gpu_reorder.hpp" #if !defined(USE_CUDA) && !defined(USE_HIP) && !defined(USE_DPCPP) namespace tamm { @@ -138,42 +139,22 @@ template void assign_gpu(gpuStream_t& thandle, T*& dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign) { - const int ndim = sdims.size(); + const size_t ndim = sdims.size(); + EXPECTS(ddims.size() == ndim && dlabels.size() == ndim && slabels.size() == ndim); + EXPECTS(dst != nullptr && (ndim == 0 || src != nullptr)); + EXPECTS(dst != src); // out-of-place only; in-place permutes need a temp buffer const Size ssize = std::accumulate(sdims.begin(), sdims.end(), Size{1}, std::multiplies()); - if(ndim <= 1 || ssize.value() == 1) { - // device-->device copy - gpuMemcpyAsync(dst, src, ssize.value(), gpuMemcpyDeviceToDevice, thandle); - } - - std::vector r_sdims; - std::transform(std::begin(sdims), std::end(sdims), std::back_inserter(r_sdims), - [](tamm::Size i) -> int { return i.value(); }); - - tamm::IntLabelVec r_dlabels = dlabels; - tamm::IntLabelVec r_slabels = slabels; - - // if(is_assign) - std::reverse(r_sdims.begin(), r_sdims.end()); - std::reverse(r_slabels.begin(), r_slabels.end()); - std::reverse(r_dlabels.begin(), r_dlabels.end()); - - int perm[ndim]; - int size[ndim]; - // T beta = is_assign ? 0 : 1; - - for(size_t i = 0; i < r_sdims.size(); i++) { size[i] = r_sdims[i]; } - for(size_t i = 0; i < r_dlabels.size(); i++) { - auto it = std::find(r_slabels.begin(), r_slabels.end(), r_dlabels[i]); - EXPECTS(it != r_slabels.end()); - perm[i] = it - r_slabels.begin(); - } - - // Out-of-place N-dimensional axis-permuting transpose (in-house reorder - // kernel; replaces the previous librett-based implementation). - int outDims[ndim]; - for(int i = 0; i < ndim; i++) { outDims[i] = size[perm[i]]; } - gpu::transpose_reorder(dst, src, ndim, outDims, perm, thandle); + if(ssize.value() == 0) return; + EXPECTS(ndim <= static_cast(gpu::reorder_maxrank)); + + // Fortran-order (reversed) spec for the column-major reorder kernel; this + // is exactly the row-major transpose the CPU assign() computes. + // Allocation-free (stack spec): no heap traffic in this hot path. + gpu::ReorderSpec spec{}; + gpu::build_reorder_spec(sdims, slabels, dlabels, spec); + gpu::transpose_reorder(dst, src, static_cast(ndim), spec.outDims, spec.perm, scale, + !is_assign, thandle); } #endif @@ -201,11 +182,11 @@ bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, assign_gpu(thandle, binter_buf_dev, binter_dims, binter_labels, T3{1}, binter_dev_in.data(), bdims, blabels, true); - // The H2D copies above and librettExecute() inside assign_gpu() are enqueued on + // The H2D copies above and transpose_reorder() inside assign_gpu() are enqueued on // `thandle` and return immediately on CUDA/HIP, but returning these staging buffers to // the pool makes their addresses instantly re-allocatable (the pool's deallocate is // host-side bookkeeping, not a stream-ordered free). Without this sync the next - // allocation can alias memory that librett is still reading. + // allocation can alias memory that the transpose is still reading. gpuStreamSynchronize(thandle); free_device_buffer(hw, ainter_dev_in); @@ -227,8 +208,17 @@ void transpose_output(ExecutionHW hw, gpuStream_t& thandle, bool gpu_trans, T1* T1*& cinter_tmp_buf_dev, bool is_assign) { #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) if(hw == ExecutionHW::GPU) { + // NOTE: always overwrite (is_assign=true) on the GPU path, matching the + // old librett assign_gpu semantics (which ignored is_assign/scale). + // Accumulation across reduction iterations here is performed by GEMM + // itself (beta=cscale=1 accumulates in cinter_tmp_buf_dev in inter + // ordering), so the output transpose must overwrite cinter_buf_dev with + // the running total each iteration. Honoring is_assign=false here (D+=T) + // double-counts: iter2 adds T1+S2 on top of D=T1, giving 2*S1+S2. + // (CPU path differs: its per-call cinter_buf is freshly zeroed, so the + // transpose is the accumulator there and must honor is_assign.) assign_gpu(thandle, cinter_buf_dev, cdims, clabels, T1{1}, cinter_tmp_buf_dev, cinter_dims, - cinter_labels, is_assign); + cinter_labels, true); return; } #endif @@ -623,7 +613,7 @@ void block_multiply( is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gpu::axpy and the librett transpose in transpose_output are stream-async on + // gpu::axpy and the reorder transpose in transpose_output are stream-async on // CUDA/HIP; sync before returning this device block to the pool, which would // otherwise make it immediately re-allocatable while that work is still reading it. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -646,16 +636,18 @@ void block_multiply( // cinter_*_dev buffers) and every device buffer allocated internally is free of // in-flight work by the time this function returns. // - // On CUDA/HIP the cuBLAS/rocBLAS gemm above is enqueued on `thandle` and returns - // immediately, whereas the memory pool's deallocate() is pure host-side bookkeeping that - // makes a block instantly re-allocatable -- it does not defer the reuse behind a stream - // event the way upstream RMM's stream-ordered free lists do. Callers that free a device - // buffer right after this call (e.g. the reduction loop in MultOp::execute_bufacc, and - // the equivalent loops in exachem's cd_ccsd_{cs,os}_ann.cpp) would otherwise hand the - // same address to the next iteration's allocation while the gemm is still reading it. - // - // DPC++ already synchronises inside gpu::gemm (gemm_event.wait()), so this is a no-op - // there; the cost on CUDA/HIP is the sync the callers' correctness already assumed. + // On all backends the gemm/transpose work above is enqueued on `thandle` and + // returns immediately (CUDA/HIP: async BLAS; DPC++: in-order queue, no waits + // inside gpu::gemm by design), whereas the memory pool's deallocate() is pure + // host-side bookkeeping that makes a block instantly re-allocatable -- it does + // not defer the reuse behind a stream event the way upstream RMM's + // stream-ordered free lists do. Callers that free a device buffer right after + // this call (e.g. the reduction loop in MultOp::execute_bufacc, and the + // equivalent loops in exachem's cd_ccsd_{cs,os}_ann.cpp) would otherwise hand + // the same address to the next iteration's allocation while the gemm is still + // reading it. This device-wide drain per block op is the price of the + // immediate-free pool; removing it requires moving the device pool to a + // stream-ordered resource (see mr/stream_ordered_memory_resource.hpp). if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } #endif diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index ffce350a3..32d0c5009 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) #include #endif @@ -26,16 +28,20 @@ template void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); -// Out-of-place N-dimensional axis-permuting transpose (reorder). -// Replaces the previous librett-based implementation. -// out : destination buffer (device), sized as the permuted tensor -// in : source buffer (device) -// ndim : tensor rank (supports up to tamm::maxrank) -// outDims : extents of the output (permuted) tensor, length ndim -// perm : output-axis -> source-axis map, length ndim +// Out-of-place N-dimensional axis-permuting transpose (in-house reorder +// kernel; fully replaces the old librett dependency). +// out : destination buffer (device), sized as the permuted tensor +// in : source buffer (device); must not alias out +// ndim : tensor rank (0 .. 8) +// outDims : extents of the output (permuted) tensor, length ndim +// perm : output-axis -> source-axis map, length ndim +// scale : every output element is multiplied by scale +// accumulate : false -> out = scale*in, true -> out += scale*in +// Enqueued on handle.first (cuda/hip stream, in-order SYCL queue); async, +// no internal synchronization. template -void transpose_reorder(T* out, const T* in, int ndim, const int* outDims, const int* perm, - gpuStream_t& handle); +void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, + T scale, bool accumulate, gpuStream_t& handle); } // namespace gpu #endif diff --git a/tests/tamm/Test_TransposeReorder.cpp b/tests/tamm/Test_TransposeReorder.cpp new file mode 100644 index 000000000..1c7ecb242 --- /dev/null +++ b/tests/tamm/Test_TransposeReorder.cpp @@ -0,0 +1,239 @@ +// Unit tests for the in-house GPU tensor transpose (reorder). +// +// Exercises the backend-agnostic logic in tamm/kernels/gpu_reorder.hpp +// (spec building, metadata, index decode, scale/accumulate apply) directly on +// the host, so no GPU is needed. Every test compares the exact helpers the +// CUDA/HIP/SYCL kernels call against an independent row-major reference for +// all permutations of ranks 0..5 (plus rank-6/8 spot checks), several shapes +// (including unit extents), and double/complex values with a range of +// scale/accumulate settings. + +#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN +#include "doctest/doctest.h" + +#include + +#include +#include +#include +#include + +using namespace tamm; +using namespace tamm::kernels::gpu; + +namespace { + +// Row-major (last-axis-fastest) strides. +std::vector row_strides(const std::vector& dims) { + const size_t n = dims.size(); + std::vector st(n, 1); + for(size_t k = n; k-- > 1;) st[k - 1] = st[k] * dims[k]; + return st; +} + +// Independent reference: dst[d] = (acc ? dst[d] + : 0) + scale*src[s] with a +// plain row-major transpose driven by the natural (non-reversed) permutation. +template +void reference_transpose(const std::vector& src, std::vector& dst, + const std::vector& sdims, const std::vector& slabels, + const std::vector& ddims, const std::vector& dlabels, T scale, + bool accum) { + const size_t ndim = sdims.size(); + const auto sst = row_strides(sdims); + const auto dst_ = row_strides(ddims); + // dest-axis -> src-axis in natural order + std::vector perm(ndim); + for(size_t i = 0; i < ndim; ++i) { + perm[i] = static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - + slabels.begin()); + } + size_t total = 1; + for(auto d: ddims) total *= d; + REQUIRE(dst.size() == total); + std::vector coord(ndim); + for(size_t lin = 0; lin < total; ++lin) { + size_t rem = lin; + for(size_t k = ndim; k-- > 0;) { + coord[k] = (ndim == 0) ? 0 : rem % ddims[k]; + if(ndim > 0) rem /= ddims[k]; + } + size_t s = 0; + for(size_t k = 0; k < ndim; ++k) s += coord[k] * sst[static_cast(perm[k])]; + dst[lin] = accum ? dst[lin] + scale * src[s] : scale * src[s]; + } +} + +// Replica of the device loop in gpu_reorder.cpp, but on the host, using the +// shipped helpers (build_reorder_spec / reorder_build_meta / +// reorder_src_index / reorder_scaled / reorder_add). This is the exact math +// each backend kernel executes; only the thread loop itself is replicated. +template +void device_math_replica(const std::vector& src, std::vector& dst, const SizeVec& sdims, + const IntLabelVec& slabels, const IntLabelVec& dlabels, T scale, + bool accum) { + const int ndim = static_cast(sdims.size()); + ReorderSpec spec{}; + build_reorder_spec(sdims, slabels, dlabels, spec); + const ReorderMeta meta = reorder_build_meta(ndim, spec.outDims, spec.perm); + const size_t total = reorder_total(ndim, spec.outDims); + REQUIRE(dst.size() == total); + double sre, sim; + if constexpr(reorder_is_complex_v) { + sre = static_cast(scale.real()); + sim = static_cast(scale.imag()); + } + else { + sre = static_cast(scale); + sim = 0.0; + } + for(size_t t = 0; t < total; ++t) { + const Idx tid = static_cast(t); + const size_t s = reorder_src_index(tid, meta); + const T y = reorder_scaled(src[s], sre, sim); + dst[t] = accum ? reorder_add(dst[t], y) : y; + } +} + +template +T make_val(size_t i) { + if constexpr(reorder_is_complex_v) { + return T(static_cast(0.25 * (i + 1)), + static_cast(-0.125 * (i + 3))); + } + else { return static_cast(0.25 * (i + 1)); } +} + +template +void check_case(const std::vector& sdims, const std::vector& slabels, + const std::vector& dlabels, T scale, bool accum) { + const size_t ndim = sdims.size(); + SizeVec sv; + for(auto d: sdims) sv.emplace_back(d); + IntLabelVec sl(slabels.begin(), slabels.end()), dl(dlabels.begin(), dlabels.end()); + + std::vector ddims(ndim); + for(size_t i = 0; i < ndim; ++i) { + const int j = + static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - slabels.begin()); + ddims[i] = sdims[static_cast(j)]; + } + size_t total = 1; + for(auto d: ddims) total *= d; + + std::vector src(total), ref(total), got32(total), got64(total), init(total); + for(size_t i = 0; i < total; ++i) { + src[i] = make_val(i + 1); + init[i] = make_val(1000 + i); + ref[i] = init[i]; + got32[i] = init[i]; + got64[i] = init[i]; + } + reference_transpose(src, ref, sdims, slabels, ddims, dlabels, scale, accum); + device_math_replica(src, got32, sv, sl, dl, scale, accum); + device_math_replica(src, got64, sv, sl, dl, scale, accum); + for(size_t i = 0; i < total; ++i) { + CHECK(got32[i] == ref[i]); + CHECK(got64[i] == ref[i]); + } + // The 64-bit total path must agree with the 32-bit path elementwise. + for(size_t i = 0; i < total; ++i) { CHECK(got64[i] == got32[i]); } +} + +template +void check_all_perms(const std::vector& sdims, T scale, bool accum) { + const size_t ndim = sdims.size(); + std::vector labels(ndim); + std::iota(labels.begin(), labels.end(), 0); + if(ndim == 0) { + check_case({}, {}, {}, scale, accum); + return; + } + std::vector perm = labels; + do { check_case(sdims, labels, perm, scale, accum); } while( + std::next_permutation(perm.begin(), perm.end())); +} + +} // namespace + +TEST_CASE("reorder rank 0..2, all perms, scales, assign/accumulate") { + using C = std::complex; + for(bool accum: {false, true}) { + for(double s: {1.0, 2.5, 0.0, -1.0}) { + check_all_perms({}, static_cast(s), accum); + check_all_perms({1}, static_cast(s), accum); + check_all_perms({5}, static_cast(s), accum); + check_all_perms({1, 1}, static_cast(s), accum); + check_all_perms({2, 3}, static_cast(s), accum); + check_all_perms({4, 1}, static_cast(s), accum); + check_all_perms({2, 3}, C(s, 0.5 * s), accum); + } + } +} + +TEST_CASE("reorder rank 3, all perms") { + using C = std::complex; + for(bool accum: {false, true}) { + check_all_perms({2, 3, 4}, 1.0, accum); + check_all_perms({1, 5, 2}, 2.0, accum); + check_all_perms({3, 1, 1}, -1.5, accum); + check_all_perms({2, 3, 4}, C(1.0, -1.0), accum); + check_all_perms({1, 1, 6}, C(0.0, 2.0), accum); + } +} + +TEST_CASE("reorder rank 4, all perms") { + using C = std::complex; + for(bool accum: {false, true}) { + check_all_perms({2, 3, 2, 4}, 1.0, accum); + check_all_perms({4, 1, 3, 2}, 0.5, accum); + check_all_perms({2, 2, 3, 2}, C(2.0, 1.0), accum); + } +} + +TEST_CASE("reorder rank 5, all perms") { + for(bool accum: {false, true}) { + check_all_perms({2, 1, 3, 2, 2}, 1.0, accum); + check_all_perms>({2, 1, 2, 2, 3}, std::complex(1.0, 1.0), + accum); + } +} + +TEST_CASE("reorder rank 6 and 8 spot checks") { + // 720 perms at rank 6: check identity, reversal, and a rotation plus a few + // random perms rather than all of them. + const std::vector d6{2, 3, 1, 4, 2, 2}; + const std::vector id6{0, 1, 2, 3, 4, 5}; + check_case(d6, id6, {5, 4, 3, 2, 1, 0}, 1.0, false); + check_case(d6, id6, {5, 4, 3, 2, 1, 0}, 3.0, true); + check_case(d6, id6, id6, 1.0, false); + check_case(d6, id6, {1, 2, 3, 4, 5, 0}, -2.0, true); + check_case(d6, id6, {2, 0, 4, 1, 5, 3}, 1.0, false); + check_case>(d6, id6, {5, 4, 3, 2, 1, 0}, std::complex(0.5, -0.5), + true); + const std::vector d8{2, 1, 2, 3, 1, 2, 2, 2}; + const std::vector id8{0, 1, 2, 3, 4, 5, 6, 7}; + check_case(d8, id8, {7, 6, 5, 4, 3, 2, 1, 0}, 1.0, false); + check_case(d8, id8, id8, 1.0, false); + check_case(d8, id8, {3, 7, 1, 5, 0, 6, 2, 4}, 2.0, true); +} + +TEST_CASE("reorder metadata invariants") { + // Column-major strides + perm round-trip on a known case. + const size_t outDims[3] = {4, 2, 3}; + const int perm[3] = {2, 0, 1}; + const ReorderMeta meta = reorder_build_meta(3, outDims, perm); + CHECK(meta.outStrides[0] == 1); + CHECK(meta.outStrides[1] == 4); + CHECK(meta.outStrides[2] == 8); + // inDims = {2, 3, 4} -> inStrides {1, 2, 6} + CHECK(meta.inStrides[0] == 1); + CHECK(meta.inStrides[1] == 2); + CHECK(meta.inStrides[2] == 6); + CHECK(reorder_total(3, outDims) == 24); + CHECK_FALSE(reorder_is_identity(meta)); + const int idp[3] = {0, 1, 2}; + const ReorderMeta idm = reorder_build_meta(3, outDims, idp); + CHECK(reorder_is_identity(idm)); + CHECK(reorder_meta_fits32(meta, 24)); + CHECK_FALSE(reorder_meta_fits32(meta, size_t{1} << 33)); +} diff --git a/tests/tamm/test_tamm.cmake b/tests/tamm/test_tamm.cmake index 1a418436e..7f58e0de4 100644 --- a/tests/tamm/test_tamm.cmake +++ b/tests/tamm/test_tamm.cmake @@ -27,6 +27,7 @@ add_cxx_unit_test(Test_StrongNum) add_cxx_unit_test(Test_BoundVec) add_cxx_unit_test(Test_BlockSpan) add_cxx_unit_test(Test_MemoryPool) +add_cxx_unit_test(Test_TransposeReorder) # add_mpi_unit_test(Test_ViewTensor 2 "") # add_mpi_unit_test(Test_QR 2 "") From 4e0e79deb6db46cfa28eacb1afa88bde4e230944 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Mon, 21 Sep 2026 13:59:37 -0500 Subject: [PATCH 04/14] Replace HPTT with in-house CPU transpose Host row-major transpose_reorder_cpu (general alpha/beta, odometer loop, no per-element division) reusing the gpu_reorder math. Rewires assign, BlockAssignPlan and drops the HPTT dependency. Unit tests: 117522 assertions green; h2o/butanol2 CPU exact. --- CMakeLists.txt | 2 +- docs/install.rst | 1 - src/tamm/CMakeLists.txt | 3 +- src/tamm/block_assign_plan.hpp | 35 +++--- src/tamm/ip_hptt.hpp | 60 --------- src/tamm/ip_reorder.hpp | 45 +++++++ src/tamm/kernels/assign.hpp | 43 +++---- src/tamm/kernels/cpu_reorder.hpp | 140 +++++++++++++++++++++ src/tamm/scheduler.hpp | 2 +- tests/tamm/Test_TransposeReorder.cpp | 174 +++++++++++++++++++++++++++ 10 files changed, 397 insertions(+), 108 deletions(-) delete mode 100644 src/tamm/ip_hptt.hpp create mode 100644 src/tamm/ip_reorder.hpp create mode 100644 src/tamm/kernels/cpu_reorder.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 93265cad2..76a283e53 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -37,7 +37,7 @@ endif() set(CMSB_PROJECTS TAMM) -set(TAMM_DEPENDENCIES NJSON DOCTEST SPDLOG Eigen3 HPTT) +set(TAMM_DEPENDENCIES NJSON DOCTEST SPDLOG Eigen3) if(NOT DEFINED USE_HDF5 OR USE_HDF5) list(APPEND TAMM_DEPENDENCIES HDF5) diff --git a/docs/install.rst b/docs/install.rst index c200430ed..e4707296a 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -17,7 +17,6 @@ Dependencies **The remaining dependencies are automatically built and do not need to be installed explicitly:** * GlobalArrays -* HPTT * HDF5 * BLAS/LAPACK (BLIS and netlib-lapack are automatically built if no vendor BLAS libraries are provided) * BLAS++ and LAPACK++ diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index af76e2e78..9b4349b78 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -85,6 +85,7 @@ set(TAMM_INCLUDES kernels/multiply.hpp kernels/tamm_blas.hpp kernels/gpu_reorder.hpp + kernels/cpu_reorder.hpp op_dag.hpp interfaces.hpp op_profiler.hpp @@ -101,7 +102,7 @@ set(TAMM_INCLUDES blockops_cpu.hpp perm.hpp iteration.hpp - ip_hptt.hpp + ip_reorder.hpp blockops_blas.hpp tensor_variant.hpp op_executor.hpp diff --git a/src/tamm/block_assign_plan.hpp b/src/tamm/block_assign_plan.hpp index 6732543b1..0fce6618c 100644 --- a/src/tamm/block_assign_plan.hpp +++ b/src/tamm/block_assign_plan.hpp @@ -6,7 +6,7 @@ #include "tamm/blockops_blas.hpp" #include "tamm/blockops_cpu.hpp" #include "tamm/errors.hpp" -#include "tamm/ip_hptt.hpp" +#include "tamm/ip_reorder.hpp" #include "tamm/types.hpp" /////////////////////////////////////////////////////////////////////////////// @@ -19,9 +19,6 @@ namespace tamm { /** * @brief Block assign plan for block assignment on Tensor Ops * - * @todo HPTT calls should be fixed for mixed precision (T1 && T2) when we allow - * those operations in make_op, currently only complex and scalar mix is allowed - * */ class BlockAssignPlan { public: @@ -37,7 +34,7 @@ class BlockAssignPlan { // rhs_labels{rhs_labels}, plan_{Plan::invalid} { prep_flat_plan(lhs_labels, rhs_labels); - if(plan_ == Plan::invalid) { prep_hptt(lhs_labels, rhs_labels); } + if(plan_ == Plan::invalid) { prep_reorder(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_index_permute_plan(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_general_loop_plan(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_general_plan(lhs_labels, rhs_labels); } @@ -56,9 +53,9 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rhs); break; case Plan::flat_update: blockops::cpu::flat_update(lhs, rhs); break; - case Plan::hptt: - blockops::hptt::index_permute_hptt(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_, - rhs.block_dims()); + case Plan::reorder: + blockops::reorder::index_permute_reorder(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_, + rhs.block_dims()); break; case Plan::index_permute_assign: blockops::cpu::index_permute_assign(lhs.buf(), rhs.buf(), ip_plan_.perm_, lhs.block_dims()); @@ -118,9 +115,9 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rscale, rhs); break; case Plan::flat_update: blockops::cpu::flat_update(lhs, rscale, rhs); break; - case Plan::hptt: - blockops::hptt::index_permute_hptt(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_, - rhs.block_dims()); + case Plan::reorder: + blockops::reorder::index_permute_reorder(lscale, lhs.buf(), rscale, rhs.buf(), + ip_plan_.perm_, rhs.block_dims()); break; case Plan::index_permute_assign: blockops::cpu::index_permute_assign(lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_, @@ -200,9 +197,9 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: NOT_ALLOWED(); break; case Plan::flat_update: blockops::cpu::flat_update(lscale, lhs, rscale, rhs); break; - case Plan::hptt: - blockops::hptt::index_permute_hptt(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_, - rhs.block_dims()); + case Plan::reorder: + blockops::reorder::index_permute_reorder(lscale, lhs.buf(), rscale, rhs.buf(), + ip_plan_.perm_, rhs.block_dims()); break; case Plan::index_permute_assign: NOT_ALLOWED(); break; case Plan::index_permute_update: @@ -268,12 +265,12 @@ class BlockAssignPlan { } template - void prep_hptt(const std::vector& lhs_labels, const std::vector& rhs_labels) { + void prep_reorder(const std::vector& lhs_labels, const std::vector& rhs_labels) { prep_index_permute_plan(lhs_labels, rhs_labels); if(plan_ == Plan::invalid) { - return; // no index permute plan, no hptt plan + return; // no index permute plan, no reorder plan } - plan_ = Plan::hptt; + plan_ = Plan::reorder; } template @@ -318,7 +315,7 @@ class BlockAssignPlan { enum class Plan { flat_assign, flat_update, - hptt, + reorder, index_permute_assign, index_permute_update, ipgen_loop_assign, @@ -346,7 +343,7 @@ class BlockAssignPlan { switch(plan) { case Plan::flat_assign: return "flat_assign"; case Plan::flat_update: return "flat_update"; - case Plan::hptt: return "hptt"; + case Plan::reorder: return "reorder"; case Plan::index_permute_assign: return "index_permute_assign"; case Plan::index_permute_update: return "index_permute_update"; case Plan::ipgen_loop_assign: return "ipgen_loop_assign"; diff --git a/src/tamm/ip_hptt.hpp b/src/tamm/ip_hptt.hpp deleted file mode 100644 index c2e1a9bd9..000000000 --- a/src/tamm/ip_hptt.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#pragma once - -#include - -#include "hptt/hptt.h" -#include "tamm/errors.hpp" -#include "tamm/types.hpp" - -namespace tamm::blockops::hptt { - -/////////////////////////////////////////////////////////////////////////////// -// -// IP using HPTT -// -/////////////////////////////////////////////////////////////////////////////// - -// template -// void index_permute_assign_hptt(T* lbuf, const T* rbuf, -// const PermVector& perm_to_dest, -// const std::vector& ddims, bool is_assign, -// T scale = 1.0) { -// const int ndim = ddims.size(); -// int perm[ndim]; -// int size[ndim]; -// T beta = is_assign ? 0 : 1; -// int numThreads = 1; -// for (size_t i = 0; i < ddims.size(); i++) { -// size[i] = ddims[i].value(); -// } - -// for (size_t i = 0; i < perm_to_dest.size(); i++) { -// perm[i] = perm_to_dest[i].value(); -// } -// // create a plan (shared_ptr) -// auto plan = hptt::create_plan(perm, ndim, scale, rbuf, size, NULL, beta, lbuf, -// NULL, hptt::ESTIMATE, numThreads, NULL, true); - -// // execute the transposition -// plan->execute(); -// } - -template -void index_permute_hptt(T lscale, T* lbuf, T rscale, const T* rbuf, const PermVector& perm_to_dest, - const std::vector& sdims) { - const int ndim = sdims.size(); - int perm[ndim]; - int size[ndim]; - int num_threads = 1; - for(size_t i = 0; i < sdims.size(); i++) { size[i] = sdims[i]; } - - for(size_t i = 0; i < perm_to_dest.size(); i++) { perm[i] = perm_to_dest[i]; } - // create a plan (shared_ptr) - auto plan = ::hptt::create_plan(perm, ndim, rscale, rbuf, size, NULL, lscale, lbuf, NULL, - ::hptt::ESTIMATE, num_threads, NULL, true); - - // execute the transposition - plan->execute(); -} - -} // namespace tamm::blockops::hptt diff --git a/src/tamm/ip_reorder.hpp b/src/tamm/ip_reorder.hpp new file mode 100644 index 000000000..caeb14143 --- /dev/null +++ b/src/tamm/ip_reorder.hpp @@ -0,0 +1,45 @@ +#pragma once + +#include +#include + +#include "tamm/errors.hpp" +#include "tamm/kernels/cpu_reorder.hpp" +#include "tamm/types.hpp" + +namespace tamm::blockops::reorder { + +/////////////////////////////////////////////////////////////////////////////// +// +// Index permute using the in-house CPU reorder kernel +// (replaces the old HPTT-based index_permute_hptt) +// +/////////////////////////////////////////////////////////////////////////////// + +// Out-of-place transpose with HPTT-identical semantics: +// lbuf = lscale * lbuf + rscale * transpose(rbuf) +// perm_to_dest maps each destination axis to its source axis and sdims holds +// the source extents, both in natural order; the output extents follow as +// outDims[k] == sdims[perm_to_dest[k]]. Scale types are independent template +// parameters so call sites may mix e.g. an int literal rscale with a +// complex buffer (as the old single-T HPTT wrapper allowed via conversion). +template +void index_permute_reorder(BL lscale, T1* lbuf, BA rscale, const T2* rbuf, + const PermVector& perm_to_dest, const std::vector& sdims) { + const size_t ndim = sdims.size(); + EXPECTS(perm_to_dest.size() == ndim); + EXPECTS(ndim <= static_cast(kernels::gpu::reorder_maxrank)); + EXPECTS(lbuf != nullptr && (ndim == 0 || rbuf != nullptr)); + + size_t outDims[kernels::gpu::reorder_maxrank] = {}; + int perm[kernels::gpu::reorder_maxrank] = {}; + for(size_t i = 0; i < ndim; ++i) { + perm[i] = static_cast(perm_to_dest[i]); + EXPECTS(perm[i] >= 0 && static_cast(perm[i]) < ndim); + outDims[i] = sdims[static_cast(perm[i])]; + } + kernels::cpu::transpose_reorder_cpu(lbuf, rbuf, static_cast(ndim), outDims, perm, rscale, + lscale); +} + +} // namespace tamm::blockops::reorder diff --git a/src/tamm/kernels/assign.hpp b/src/tamm/kernels/assign.hpp index 33fa23f2a..c34dea08d 100644 --- a/src/tamm/kernels/assign.hpp +++ b/src/tamm/kernels/assign.hpp @@ -1,11 +1,10 @@ #pragma once #include "tamm/errors.hpp" +#include "tamm/kernels/cpu_reorder.hpp" #include "tamm/types.hpp" #include "tamm/utils.hpp" -#include "hptt/hptt.h" - #include #include #include @@ -374,31 +373,25 @@ void ip_gen_loop(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T sca } template -void ip_hptt(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, - const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign = true) { - // EXPECTS(ddims.size() == dlabels.size()); - // EXPECTS(sdims.size() == slabels.size()); - // EXPECTS(ddims.size() == sdims.size()); - // EXPECTS(src != nullptr); - // EXPECTS(dst != nullptr); - - const int ndim = ddims.size(); - int perm[ndim]; - int size[ndim]; - T beta = is_assign ? 0 : 1; - int numThreads = 1; - for(size_t i = 0; i < sdims.size(); i++) { size[i] = sdims[i].value(); } - for(size_t i = 0; i < dlabels.size(); i++) { +void ip_reorder(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, + const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign = true) { + const size_t ndim = ddims.size(); + EXPECTS(sdims.size() == ndim && dlabels.size() == ndim && slabels.size() == ndim); + EXPECTS(ndim <= static_cast(kernels::gpu::reorder_maxrank)); + + // Natural-order output extents + output-axis -> source-axis map, exactly + // the (size, perm) pair the old HPTT plan consumed. + size_t outDims[kernels::gpu::reorder_maxrank] = {}; + int perm[kernels::gpu::reorder_maxrank] = {}; + for(size_t i = 0; i < ndim; i++) { auto it = std::find(slabels.begin(), slabels.end(), dlabels[i]); EXPECTS(it != slabels.end()); - perm[i] = it - slabels.begin(); + const size_t j = static_cast(it - slabels.begin()); + perm[i] = static_cast(j); + outDims[i] = sdims[j].value(); } - // create a plan (shared_ptr) - auto plan = hptt::create_plan(perm, ndim, scale, src, size, NULL, beta, dst, NULL, hptt::ESTIMATE, - numThreads, NULL, true); - - // execute the transposition - plan->execute(); + kernels::cpu::transpose_reorder_cpu(dst, src, static_cast(ndim), outDims, perm, scale, + is_assign ? T{0} : T{1}); } /** @@ -480,7 +473,7 @@ void assign(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, c if(is_assign) { internal::index_permute(dst, src, perm_to_dest, ddims, scale); } else { internal::index_permute_acc(dst, src, perm_to_dest, ddims, scale); } } - else internal::ip_hptt(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); + else internal::ip_reorder(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } else { internal::ip_gen_loop(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } } diff --git a/src/tamm/kernels/cpu_reorder.hpp b/src/tamm/kernels/cpu_reorder.hpp new file mode 100644 index 000000000..14e02a1cf --- /dev/null +++ b/src/tamm/kernels/cpu_reorder.hpp @@ -0,0 +1,140 @@ +#pragma once + +// Host-only CPU tensor transpose, the HPTT replacement. +// +// Row-major (TAMM-native) counterpart of the GPU kernel in gpu_reorder.hpp. +// Reuses its ReorderMeta storage, elementwise scale/add helpers and identity +// check, so both backends implement the same math; only the metadata build +// differs (row-major strides, natural axis order). +// +// out = beta * out + alpha * transpose(in) +// +// beta==0 never reads out; beta==1 accumulates. Single-threaded and +// allocation-free except on the aliased in-place path. + +#include "tamm/kernels/gpu_reorder.hpp" + +#include +#include +#include +#include + +namespace tamm::kernels::cpu { + +// Split a host-side alpha/beta into raw {real, imag} doubles. +template +inline void reorder_split_scale_host(S scale, double& re, double& im) { + if constexpr(gpu::reorder_is_complex_v) { + re = static_cast(scale.real()); + im = static_cast(scale.imag()); + } + else { + re = static_cast(scale); + im = 0.0; + } +} + +template +inline bool reorder_scale_is_zero_host(S scale) { + if constexpr(gpu::reorder_is_complex_v) { return scale.real() == 0 && scale.imag() == 0; } + else { return scale == S{0}; } +} + +template +inline bool reorder_scale_is_one_host(S scale) { + if constexpr(gpu::reorder_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } + else { return scale == S{1}; } +} + +// Build row-major metadata from natural output extents + dst->src perm. +// Source extents follow from inDims[perm[k]] == outDims[k]. +inline gpu::ReorderMeta reorder_build_meta_rowmajor(int ndim, const size_t* outDims, + const int* perm) { + EXPECTS(ndim >= 0 && ndim <= gpu::reorder_maxrank); + gpu::ReorderMeta meta{}; + meta.ndim = ndim; + if(ndim == 0) return meta; + + size_t inDims[gpu::reorder_maxrank] = {}; + for(int k = 0; k < ndim; ++k) { + EXPECTS(perm[k] >= 0 && perm[k] < ndim); + meta.outDims[k] = outDims[k]; + meta.perm[k] = perm[k]; + inDims[perm[k]] = outDims[k]; + } + meta.outStrides[ndim - 1] = 1; + meta.inStrides[ndim - 1] = 1; + for(int k = ndim - 2; k >= 0; --k) { + meta.outStrides[k] = meta.outStrides[k + 1] * outDims[k + 1]; + meta.inStrides[k] = meta.inStrides[k + 1] * inDims[k + 1]; + } + return meta; +} + +template +inline void transpose_reorder_cpu_impl(T* out, const T* in, const gpu::ReorderMeta& meta, + size_t total, double alpha_re, double alpha_im, B beta) { + const bool beta_is_zero = reorder_scale_is_zero_host(beta); + const bool beta_is_one = reorder_scale_is_one_host(beta); + double beta_re = 0.0, beta_im = 0.0; + if(!beta_is_zero && !beta_is_one) reorder_split_scale_host(beta, beta_re, beta_im); + + // Odometer over the output coordinates: no division/modulo per element, + // just ndim multiply-adds plus an amortized single increment. + const int ndim = meta.ndim; + size_t coord[gpu::reorder_maxrank] = {}; + for(size_t t = 0; t < total; ++t) { + size_t src = 0; + for(int k = 0; k < ndim; ++k) src += coord[k] * meta.inStrides[meta.perm[k]]; + const T y = gpu::reorder_scaled(in[src], alpha_re, alpha_im); + if(beta_is_zero) { out[t] = y; } + else if(beta_is_one) { out[t] = gpu::reorder_add(out[t], y); } + else { out[t] = gpu::reorder_add(gpu::reorder_scaled(out[t], beta_re, beta_im), y); } + for(int k = ndim - 1; k >= 0; --k) { + if(++coord[k] < meta.outDims[k]) break; + coord[k] = 0; + } + } +} + +// out = beta * out + alpha * transpose(in); outDims/perm in natural order. +template +inline void transpose_reorder_cpu(T* out, const T* in, int ndim, const size_t* outDims, + const int* perm, A alpha, B beta) { + EXPECTS(ndim >= 0 && ndim <= gpu::reorder_maxrank); + if(ndim == 0) { + double alpha_re, alpha_im, beta_re, beta_im; + reorder_split_scale_host(alpha, alpha_re, alpha_im); + reorder_split_scale_host(beta, beta_re, beta_im); + const T y = gpu::reorder_scaled(in[0], alpha_re, alpha_im); + if(reorder_scale_is_zero_host(beta)) { out[0] = y; } + else { out[0] = gpu::reorder_add(gpu::reorder_scaled(out[0], beta_re, beta_im), y); } + return; + } + EXPECTS(outDims != nullptr && perm != nullptr && in != nullptr && out != nullptr); + + const gpu::ReorderMeta meta = reorder_build_meta_rowmajor(ndim, outDims, perm); + const size_t total = gpu::reorder_total(ndim, outDims); + if(total == 0) return; + + double alpha_re, alpha_im; + reorder_split_scale_host(alpha, alpha_re, alpha_im); + + // Identity + alpha==1 + beta==0: pure copy. + if(reorder_scale_is_zero_host(beta) && reorder_scale_is_one_host(alpha) && + gpu::reorder_is_identity(meta)) { + std::memcpy(out, in, total * sizeof(T)); + return; + } + + // Aliased non-identity permute goes via a temp buffer. + if(out == in && !gpu::reorder_is_identity(meta)) { + std::vector tmp(in, in + total); + transpose_reorder_cpu_impl(out, tmp.data(), meta, total, alpha_re, alpha_im, beta); + return; + } + + transpose_reorder_cpu_impl(out, in, meta, total, alpha_re, alpha_im, beta); +} + +} // namespace tamm::kernels::cpu diff --git a/src/tamm/scheduler.hpp b/src/tamm/scheduler.hpp index a3cf6720e..bad59bf90 100644 --- a/src/tamm/scheduler.hpp +++ b/src/tamm/scheduler.hpp @@ -6,7 +6,7 @@ #include "ga/ga-mpi.h" #include "tamm/dag_impl.hpp" #include "tamm/execution_context.hpp" -#include "tamm/ip_hptt.hpp" +#include "tamm/ip_reorder.hpp" #include "tamm/ops.hpp" #include "tamm/tensor.hpp" diff --git a/tests/tamm/Test_TransposeReorder.cpp b/tests/tamm/Test_TransposeReorder.cpp index 1c7ecb242..0a991739e 100644 --- a/tests/tamm/Test_TransposeReorder.cpp +++ b/tests/tamm/Test_TransposeReorder.cpp @@ -11,6 +11,9 @@ #define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN #include "doctest/doctest.h" +#include +#include +#include #include #include @@ -237,3 +240,174 @@ TEST_CASE("reorder metadata invariants") { CHECK(reorder_meta_fits32(meta, 24)); CHECK_FALSE(reorder_meta_fits32(meta, size_t{1} << 33)); } + +// --------------------------------------------------------------------------- +// CPU reorder kernel (HPTT replacement): same natural-order convention as the +// reference above, so no reversal is needed. Covers general alpha/beta +// (HPTT supported arbitrary output scaling), both entry points +// (kernels::internal::ip_reorder and blockops::reorder::index_permute_reorder). +// --------------------------------------------------------------------------- + +using namespace tamm::kernels::cpu; + +// Independent reference with general output scaling: +// dst[d] = beta*dst[d] + alpha*src[s], plain row-major transpose. +template +void reference_transpose_beta(const std::vector& src, std::vector& dst, + const std::vector& sdims, const std::vector& slabels, + const std::vector& ddims, const std::vector& dlabels, + T alpha, T beta) { + const size_t ndim = sdims.size(); + const auto sst = row_strides(sdims); + const auto dst_ = row_strides(ddims); + std::vector perm(ndim); + for(size_t i = 0; i < ndim; ++i) { + perm[i] = static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - + slabels.begin()); + } + size_t total = 1; + for(auto d: ddims) total *= d; + REQUIRE(dst.size() == total); + std::vector coord(ndim); + for(size_t lin = 0; lin < total; ++lin) { + size_t rem = lin; + for(size_t k = ndim; k-- > 0;) { + coord[k] = (ndim == 0) ? 0 : rem % ddims[k]; + if(ndim > 0) rem /= ddims[k]; + } + size_t s = 0; + for(size_t k = 0; k < ndim; ++k) s += coord[k] * sst[static_cast(perm[k])]; + dst[lin] = beta * dst[lin] + alpha * src[s]; + } +} + +template +void check_cpu_case(const std::vector& sdims, const std::vector& slabels, + const std::vector& dlabels, T alpha, T beta) { + const size_t ndim = sdims.size(); + std::vector ddims(ndim); + for(size_t i = 0; i < ndim; ++i) { + const int j = + static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - slabels.begin()); + ddims[i] = sdims[static_cast(j)]; + } + size_t total = 1; + for(auto d: ddims) total *= d; + + std::vector src(total), ref(total), got(total), init(total); + for(size_t i = 0; i < total; ++i) { + src[i] = make_val(i + 1); + init[i] = make_val(1000 + i); + ref[i] = init[i]; + got[i] = init[i]; + } + reference_transpose_beta(src, ref, sdims, slabels, ddims, dlabels, alpha, beta); + + // 1. Raw kernel entry point (natural-order outDims + dst->src perm). + { + std::vector perm(ndim); + for(size_t i = 0; i < ndim; ++i) { + perm[i] = static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - + slabels.begin()); + } + std::copy(init.begin(), init.end(), got.begin()); + transpose_reorder_cpu(got.data(), src.data(), static_cast(ndim), ddims.data(), + perm.data(), alpha, beta); + for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); + } + + // 2. assign.hpp path (SizeVec/IntLabelVec, is_assign <-> beta 0/1 only). + if(beta == T{0} || beta == T{1}) { + SizeVec sv; + for(auto d: sdims) sv.emplace_back(d); + IntLabelVec sl(slabels.begin(), slabels.end()), dl(dlabels.begin(), dlabels.end()); + SizeVec dv; + for(auto d: ddims) dv.emplace_back(d); + std::copy(init.begin(), init.end(), got.begin()); + tamm::internal::ip_reorder(got.data(), dv, dl, alpha, src.data(), sv, sl, + beta == T{0}); + for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); + } + + // 3. block_assign path (PermVector + source block_dims). + { + PermVector perm; + for(size_t i = 0; i < ndim; ++i) { + const int j = static_cast(std::find(slabels.begin(), slabels.end(), dlabels[i]) - + slabels.begin()); + perm.push_back(static_cast(j)); + } + std::copy(init.begin(), init.end(), got.begin()); + tamm::blockops::reorder::index_permute_reorder(beta, got.data(), alpha, src.data(), perm, + sdims); + for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); + } +} + +template +void check_cpu_all_perms(const std::vector& sdims, T alpha, T beta) { + const size_t ndim = sdims.size(); + std::vector labels(ndim); + std::iota(labels.begin(), labels.end(), 0); + if(ndim == 0) { + check_cpu_case({}, {}, {}, alpha, beta); + return; + } + std::vector perm = labels; + do { check_cpu_case(sdims, labels, perm, alpha, beta); } while( + std::next_permutation(perm.begin(), perm.end())); +} + +TEST_CASE("cpu reorder rank 0..2, all perms, general alpha/beta") { + using C = std::complex; + for(auto beta_d: {0.0, 1.0, 2.5, -0.5}) { + for(auto alpha_d: {1.0, 2.5, 0.0, -1.0}) { + check_cpu_all_perms({}, alpha_d, beta_d); + check_cpu_all_perms({1}, alpha_d, beta_d); + check_cpu_all_perms({5}, alpha_d, beta_d); + check_cpu_all_perms({1, 1}, alpha_d, beta_d); + check_cpu_all_perms({2, 3}, alpha_d, beta_d); + check_cpu_all_perms({4, 1}, alpha_d, beta_d); + check_cpu_all_perms({2, 3}, C(alpha_d, 0.5), C(beta_d, -0.25)); + } + } +} + +TEST_CASE("cpu reorder rank 3..5, all perms") { + using C = std::complex; + for(auto beta_d: {0.0, 1.0, 3.0}) { + check_cpu_all_perms({2, 3, 4}, 1.0, beta_d); + check_cpu_all_perms({1, 5, 2}, 2.0, beta_d); + check_cpu_all_perms({3, 1, 1}, -1.5, beta_d); + check_cpu_all_perms({2, 3, 4}, C(1.0, -1.0), C(beta_d, 0.5)); + check_cpu_all_perms({2, 3, 2, 4}, 0.5, beta_d); + check_cpu_all_perms({2, 2, 3, 2}, C(2.0, 1.0), C(beta_d, 0.0)); + check_cpu_all_perms({2, 1, 3, 2, 2}, 1.0, beta_d); + check_cpu_all_perms({2, 1, 2, 2, 3}, C(1.0, 1.0), C(beta_d, 1.0)); + } +} + +TEST_CASE("cpu reorder rank 6 spot checks + metadata invariants") { + const std::vector d6{2, 3, 1, 4, 2, 2}; + const std::vector id6{0, 1, 2, 3, 4, 5}; + check_cpu_case(d6, id6, {5, 4, 3, 2, 1, 0}, 1.0, 0.0); + check_cpu_case(d6, id6, {5, 4, 3, 2, 1, 0}, 3.0, 1.0); + check_cpu_case(d6, id6, id6, 1.0, 0.0); + check_cpu_case(d6, id6, {1, 2, 3, 4, 5, 0}, -2.0, 2.0); + check_cpu_case>(d6, id6, {5, 4, 3, 2, 1, 0}, + std::complex(0.5, -0.5), + std::complex(1.0, 0.0)); + // Row-major strides on a known case: outDims {4,2,3} -> {1,3,12}... strides + // count from the last axis: {2*3, 3, 1} = {6, 3, 1}. + const size_t outDims[3] = {4, 2, 3}; + const int perm[3] = {2, 0, 1}; + const ReorderMeta meta = reorder_build_meta_rowmajor(3, outDims, perm); + CHECK(meta.outStrides[0] == 6); + CHECK(meta.outStrides[1] == 3); + CHECK(meta.outStrides[2] == 1); + // inDims = {2, 3, 4} -> inStrides {12, 4, 1}. + CHECK(meta.inStrides[0] == 12); + CHECK(meta.inStrides[1] == 4); + CHECK(meta.inStrides[2] == 1); + CHECK_FALSE(reorder_is_identity(meta)); +} From 81ad68df4e2397b4c75644c61b604199b87b3b0a Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Mon, 21 Sep 2026 16:01:42 -0500 Subject: [PATCH 05/14] Streamline GPU transpose path: flat SYCL range, single end-drain per block op --- src/tamm/kernels/gpu_blas.cpp | 4 -- src/tamm/kernels/gpu_reorder.cpp | 67 ++++++++------------------------ src/tamm/kernels/multiply.hpp | 49 +++++++++++++++-------- 3 files changed, 50 insertions(+), 70 deletions(-) diff --git a/src/tamm/kernels/gpu_blas.cpp b/src/tamm/kernels/gpu_blas.cpp index 8b7513f92..7b0d27349 100644 --- a/src/tamm/kernels/gpu_blas.cpp +++ b/src/tamm/kernels/gpu_blas.cpp @@ -44,15 +44,11 @@ template void tamm::kernels::gpu::gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, const T beta, T1* C, int ldc, gpuStream_t& handle) { #if defined(USE_DPCPP) - #ifdef USE_PORT_BLAS blas::SB_Handle sb_handle(handle.first); blas::internal::_gemm(sb_handle, 'n', 'n', n, m, k, alpha, const_cast(B), ldb, const_cast(A), lda, beta, C, ldc, {}); - // No wait: the caller's in-order queue preserves ordering and - // block_multiply() synchronizes before pool buffers are recycled. #else - // No wait (same reason); the event is intentionally discarded. oneapi::mkl::blas::column_major::gemm(handle.first, oneapi::mkl::transpose::N, oneapi::mkl::transpose::N, n, m, k, alpha, B, ldb, A, lda, beta, C, ldc); diff --git a/src/tamm/kernels/gpu_reorder.cpp b/src/tamm/kernels/gpu_reorder.cpp index 8ba766c9f..1e3b7c1d5 100644 --- a/src/tamm/kernels/gpu_reorder.cpp +++ b/src/tamm/kernels/gpu_reorder.cpp @@ -2,17 +2,8 @@ #include "tamm_blas.hpp" +#include #include -#include -#include - -#if defined(USE_CUDA) -#include -#elif defined(USE_HIP) -#include -#endif -// For USE_DPCPP, SYCL is pulled in transitively via tamm_blas.hpp -> -// gpu_streams.hpp -> sycl_device.hpp. namespace tamm::kernels::gpu { @@ -56,24 +47,6 @@ reorder_kernel(T* out, const T* in, ReorderMeta meta, size_t total, double scale out[tid] = accumulate ? reorder_add(out[tid], y) : y; } } - -inline void reorder_check_launch() { -#if defined(USE_CUDA) - const cudaError_t err = cudaGetLastError(); - if(err != cudaSuccess) { - std::ostringstream msg; - msg << "TAMM reorder kernel launch failed: " << cudaGetErrorString(err); - throw std::runtime_error(msg.str()); - } -#elif defined(USE_HIP) - const hipError_t err = hipGetLastError(); - if(err != hipSuccess) { - std::ostringstream msg; - msg << "TAMM reorder kernel launch failed: " << hipGetErrorString(err); - throw std::runtime_error(msg.str()); - } -#endif -} #endif // USE_CUDA || USE_HIP template @@ -96,7 +69,6 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con meta.ndim = 0; reorder_kernel <<<1, 1, 0, handle.first>>>(out, in, meta, 1, scale_re, scale_im, accumulate); - reorder_check_launch(); } #elif defined(USE_DPCPP) if(reorder_scale_is_one(scale) && !accumulate) { @@ -127,11 +99,12 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con return; } +#if defined(USE_CUDA) || defined(USE_HIP) + // Bounded grid + stride loop (CUDA grid dims cap at 2^31-1, so a 1:1 + // mapping is impossible for huge totals). constexpr size_t block = 256; constexpr size_t maxBlocks = 1 << 20; // stride loop covers the rest const size_t nblocks = std::min((total + block - 1) / block, maxBlocks); - -#if defined(USE_CUDA) || defined(USE_HIP) if(reorder_meta_fits32(meta, total)) { reorder_kernel<<(nblocks), static_cast(block), 0, handle.first>>>(out, in, meta, total, scale_re, scale_im, @@ -142,30 +115,24 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con handle.first>>>(out, in, meta, total, scale_re, scale_im, accumulate); } - reorder_check_launch(); #elif defined(USE_DPCPP) - const sycl::nd_range<1> ndr(sycl::range<1>(nblocks * block), sycl::range<1>(block)); + // Flat one-item-per-element range: no grid sizing, no stride loop. (A + // literal single_task would serialize the whole transpose onto one + // work-item.) if(reorder_meta_fits32(meta, total)) { - handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { - const size_t tid0 = item.get_global_id(0); - const size_t stride = item.get_global_range(0); - for(size_t t = tid0; t < total; t += stride) { - const uint32_t tid = static_cast(t); - const size_t src = reorder_src_index(tid, meta); - const T y = reorder_scaled(in[src], scale_re, scale_im); - out[tid] = accumulate ? reorder_add(out[tid], y) : y; - } + handle.first.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { + const uint32_t tid = static_cast(idx[0]); + const size_t src = reorder_src_index(tid, meta); + const T y = reorder_scaled(in[src], scale_re, scale_im); + out[tid] = accumulate ? reorder_add(out[tid], y) : y; }); } else { - handle.first.parallel_for(ndr, [=](sycl::nd_item<1> item) { - const size_t tid0 = item.get_global_id(0); - const size_t stride = item.get_global_range(0); - for(size_t t = tid0; t < total; t += stride) { - const size_t src = reorder_src_index(t, meta); - const T y = reorder_scaled(in[src], scale_re, scale_im); - out[t] = accumulate ? reorder_add(out[t], y) : y; - } + handle.first.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { + const size_t t = idx[0]; + const size_t src = reorder_src_index(t, meta); + const T y = reorder_scaled(in[src], scale_re, scale_im); + out[t] = accumulate ? reorder_add(out[t], y) : y; }); } // No wait: the caller's in-order queue preserves ordering, and diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index 3220c59c3..e8e91f64f 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -5,6 +5,7 @@ #include "tamm/types.hpp" #include +#include #include #include // for std::memset #include @@ -27,6 +28,15 @@ namespace tamm { namespace kernels { +// Device staging owned by one transpose_inputs() call. Type-erased to bytes +// (the pool frees exactly size_bytes()) so block_multiply() can hold it in a +// single function-scope variable across its differently-typed branches and +// return it to the pool under the end-of-call drain. +struct GpuTransposeStaging { + std::span a{}; + std::span b{}; +}; + template void copy_data_to_gpu(ExecutionHW hw, gpuStream_t& thandle, const T2* ainter_buf, size_t asize, T2* ainter_buf_dev, const T3* binter_buf, size_t bsize, T3* binter_buf_dev) { @@ -164,7 +174,8 @@ bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, size_t asize, const SizeVec& adims, const IntLabelVec& alabels, T3* binter_buf, const SizeVec& binter_dims, const IntLabelVec& binter_labels, const T3* bbuf, size_t bsize, const SizeVec& bdims, - const IntLabelVec& blabels, T2*& ainter_buf_dev, T3*& binter_buf_dev) { + const IntLabelVec& blabels, T2*& ainter_buf_dev, T3*& binter_buf_dev, + GpuTransposeStaging& staging) { bool gpu_trans = false; #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) @@ -182,15 +193,12 @@ bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, assign_gpu(thandle, binter_buf_dev, binter_dims, binter_labels, T3{1}, binter_dev_in.data(), bdims, blabels, true); - // The H2D copies above and transpose_reorder() inside assign_gpu() are enqueued on - // `thandle` and return immediately on CUDA/HIP, but returning these staging buffers to - // the pool makes their addresses instantly re-allocatable (the pool's deallocate is - // host-side bookkeeping, not a stream-ordered free). Without this sync the next - // allocation can alias memory that the transpose is still reading. - gpuStreamSynchronize(thandle); - - free_device_buffer(hw, ainter_dev_in); - free_device_buffer(hw, binter_dev_in); + // Staging is NOT freed here: it stays alive across GEMM/output-transpose + // and is returned to the pool just before block_multiply()'s single end + // drain, saving one device-wide sync per call. (On the in-order queue all + // of the above is ordered behind prior work by construction.) + staging.a = std::as_writable_bytes(ainter_dev_in); + staging.b = std::as_writable_bytes(binter_dev_in); return gpu_trans; } @@ -356,6 +364,9 @@ void block_multiply( // int breduce_ld = B * bbatch_ld; bool gpu_trans = false; + // Input staging for the GPU path; returned to the pool under the + // end-of-call drain (see transpose_inputs). + GpuTransposeStaging transpose_staging{}; std::span cinter_span = allocate_host_buffer(hw, static_cast(csize.value())); T1* cinter_buf = cinter_span.data(); @@ -381,7 +392,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, ainter_buf_dev, - binter_buf_dev); + binter_buf_dev, transpose_staging); if(!gpu_trans) copy_data_to_gpu(hw, thandle, ainter_buf, asize.value(), ainter_buf_dev, binter_buf, @@ -422,7 +433,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf_complex, bsize.value(), bdims, blabels, - ainter_buf_dev, bbuf_complex_dev); + ainter_buf_dev, bbuf_complex_dev, transpose_staging); if(!gpu_trans) { bbuf_complex = binter_buf; @@ -458,7 +469,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf_real, bsize.value(), bdims, blabels, - ainter_buf_dev, bbuf_real_dev); + ainter_buf_dev, bbuf_real_dev, transpose_staging); if(!gpu_trans) { bbuf_real = binter_buf; @@ -505,7 +516,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf_complex, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, - blabels, abuf_complex_dev, binter_buf_dev); + blabels, abuf_complex_dev, binter_buf_dev, transpose_staging); if(!gpu_trans) { abuf_complex = ainter_buf; @@ -542,7 +553,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf_real, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, - abuf_real_dev, binter_buf_dev); + abuf_real_dev, binter_buf_dev, transpose_staging); if(!gpu_trans) { abuf_real = ainter_buf; @@ -590,7 +601,7 @@ void block_multiply( gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, - ainter_buf_dev, binter_buf_dev); + ainter_buf_dev, binter_buf_dev, transpose_staging); if(!gpu_trans) { copy_data_to_gpu(hw, thandle, ainter_buf, asize.value(), ainter_buf_dev, binter_buf, @@ -648,6 +659,12 @@ void block_multiply( // reading it. This device-wide drain per block op is the price of the // immediate-free pool; removing it requires moving the device pool to a // stream-ordered resource (see mr/stream_ordered_memory_resource.hpp). + // The transpose-input staging is returned here, covered by the same drain + // (on the in-order queue the GEMM/output-transpose above are ordered + // behind the input transposes by construction), so each call now pays a + // single device-wide sync instead of two. + free_device_buffer(hw, transpose_staging.a); + free_device_buffer(hw, transpose_staging.b); if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } #endif From e0a524c83d0a7e3e51b3b212db22bcfeb9c6b9c4 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Mon, 21 Sep 2026 17:38:42 -0500 Subject: [PATCH 06/14] Simplify rank-0 transpose path in gpu_reorder --- src/tamm/kernels/gpu_reorder.cpp | 33 ++++++++++++-------------------- 1 file changed, 12 insertions(+), 21 deletions(-) diff --git a/src/tamm/kernels/gpu_reorder.cpp b/src/tamm/kernels/gpu_reorder.cpp index 1e3b7c1d5..331e9e3e7 100644 --- a/src/tamm/kernels/gpu_reorder.cpp +++ b/src/tamm/kernels/gpu_reorder.cpp @@ -56,30 +56,24 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con EXPECTS(out != in); // out-of-place only; in-place permutes need a temp buffer if(ndim == 0) { // Rank-0: single element, still honors scale/accumulate. + if(reorder_scale_is_one(scale) && !accumulate) { + gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); + return; + } double scale_re, scale_im; reorder_split_scale(scale, scale_re, scale_im); #if defined(USE_CUDA) || defined(USE_HIP) // Route through the kernel for stream ordering (async, like every // other path through this function). - if(reorder_scale_is_one(scale) && !accumulate) { - gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); - } - else { - ReorderMeta meta{}; - meta.ndim = 0; - reorder_kernel - <<<1, 1, 0, handle.first>>>(out, in, meta, 1, scale_re, scale_im, accumulate); - } + ReorderMeta meta{}; + meta.ndim = 0; + reorder_kernel + <<<1, 1, 0, handle.first>>>(out, in, meta, 1, scale_re, scale_im, accumulate); #elif defined(USE_DPCPP) - if(reorder_scale_is_one(scale) && !accumulate) { - gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); - } - else { - handle.first.parallel_for(sycl::range<1>(1), [=](sycl::id<1>) { - const T y = reorder_scaled(in[0], scale_re, scale_im); - out[0] = accumulate ? reorder_add(out[0], y) : y; - }); - } + handle.first.parallel_for(sycl::range<1>(1), [=](sycl::id<1>) { + const T y = reorder_scaled(in[0], scale_re, scale_im); + out[0] = accumulate ? reorder_add(out[0], y) : y; + }); #endif return; } @@ -116,9 +110,6 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con accumulate); } #elif defined(USE_DPCPP) - // Flat one-item-per-element range: no grid sizing, no stride loop. (A - // literal single_task would serialize the whole transpose onto one - // work-item.) if(reorder_meta_fits32(meta, total)) { handle.first.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { const uint32_t tid = static_cast(idx[0]); From acfae716878ab72ca73019853338b989fe3b4d7f Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Tue, 22 Sep 2026 11:38:51 -0500 Subject: [PATCH 07/14] Rename transpose/reorder facility to permute --- src/tamm/CMakeLists.txt | 12 +-- src/tamm/block_assign_plan.hpp | 26 ++--- src/tamm/{ip_reorder.hpp => ip_permute.hpp} | 25 +++-- src/tamm/kernels/assign.hpp | 14 +-- .../{cpu_reorder.hpp => cpu_permute.hpp} | 78 +++++++------- .../{gpu_reorder.cpp => gpu_permute.cpp} | 72 ++++++------- .../{gpu_reorder.hpp => gpu_permute.hpp} | 100 +++++++++--------- src/tamm/kernels/multiply.hpp | 92 ++++++++-------- src/tamm/kernels/tamm_blas.hpp | 6 +- src/tamm/scheduler.hpp | 2 +- ...eReorder.cpp => Test_TransposePermute.cpp} | 80 +++++++------- tests/tamm/test_tamm.cmake | 2 +- 12 files changed, 254 insertions(+), 255 deletions(-) rename src/tamm/{ip_reorder.hpp => ip_permute.hpp} (59%) rename src/tamm/kernels/{cpu_reorder.hpp => cpu_permute.hpp} (56%) rename src/tamm/kernels/{gpu_reorder.cpp => gpu_permute.cpp} (65%) rename src/tamm/kernels/{gpu_reorder.hpp => gpu_permute.hpp} (68%) rename tests/tamm/{Test_TransposeReorder.cpp => Test_TransposePermute.cpp} (86%) diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index 9b4349b78..80a69ffe3 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -13,12 +13,12 @@ if(${PROJECT_NAME}_HAS_CUDA OR ${PROJECT_NAME}_HAS_HIP OR ${PROJECT_NAME}_HAS_DP set(tamm_gpu_SRCS kernels/gpu_blas.cpp kernels/gpu_lapack.cpp - kernels/gpu_reorder.cpp + kernels/gpu_permute.cpp ) if(${PROJECT_NAME}_HAS_CUDA) - # gpu_reorder.cpp holds __global__ kernels + <<<>>> launches, so nvcc + # gpu_permute.cpp holds __global__ kernels + <<<>>> launches, so nvcc # (not the C++ compiler) must compile it even though it is a .cpp file. - set_source_files_properties(kernels/gpu_reorder.cpp PROPERTIES LANGUAGE CUDA) + set_source_files_properties(kernels/gpu_permute.cpp PROPERTIES LANGUAGE CUDA) endif() # NOTE: hipcc compiles .cpp HIP code directly, and icpx compiles the SYCL # branch, so no LANGUAGE override is needed for HIP/DPCPP builds. @@ -84,8 +84,8 @@ set(TAMM_INCLUDES kernels/assign.hpp kernels/multiply.hpp kernels/tamm_blas.hpp - kernels/gpu_reorder.hpp - kernels/cpu_reorder.hpp + kernels/gpu_permute.hpp + kernels/cpu_permute.hpp op_dag.hpp interfaces.hpp op_profiler.hpp @@ -102,7 +102,7 @@ set(TAMM_INCLUDES blockops_cpu.hpp perm.hpp iteration.hpp - ip_reorder.hpp + ip_permute.hpp blockops_blas.hpp tensor_variant.hpp op_executor.hpp diff --git a/src/tamm/block_assign_plan.hpp b/src/tamm/block_assign_plan.hpp index 0fce6618c..5e7849767 100644 --- a/src/tamm/block_assign_plan.hpp +++ b/src/tamm/block_assign_plan.hpp @@ -6,7 +6,7 @@ #include "tamm/blockops_blas.hpp" #include "tamm/blockops_cpu.hpp" #include "tamm/errors.hpp" -#include "tamm/ip_reorder.hpp" +#include "tamm/ip_permute.hpp" #include "tamm/types.hpp" /////////////////////////////////////////////////////////////////////////////// @@ -34,7 +34,7 @@ class BlockAssignPlan { // rhs_labels{rhs_labels}, plan_{Plan::invalid} { prep_flat_plan(lhs_labels, rhs_labels); - if(plan_ == Plan::invalid) { prep_reorder(lhs_labels, rhs_labels); } + if(plan_ == Plan::invalid) { prep_permute(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_index_permute_plan(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_general_loop_plan(lhs_labels, rhs_labels); } if(plan_ == Plan::invalid) { prep_general_plan(lhs_labels, rhs_labels); } @@ -53,8 +53,8 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rhs); break; case Plan::flat_update: blockops::cpu::flat_update(lhs, rhs); break; - case Plan::reorder: - blockops::reorder::index_permute_reorder(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_, + case Plan::permute: + blockops::permute::index_permute(lscale, lhs.buf(), 1, rhs.buf(), ip_plan_.perm_, rhs.block_dims()); break; case Plan::index_permute_assign: @@ -115,8 +115,8 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: blockops::cpu::flat_assign(lhs, rscale, rhs); break; case Plan::flat_update: blockops::cpu::flat_update(lhs, rscale, rhs); break; - case Plan::reorder: - blockops::reorder::index_permute_reorder(lscale, lhs.buf(), rscale, rhs.buf(), + case Plan::permute: + blockops::permute::index_permute(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_, rhs.block_dims()); break; case Plan::index_permute_assign: @@ -197,8 +197,8 @@ class BlockAssignPlan { switch(plan_) { case Plan::flat_assign: NOT_ALLOWED(); break; case Plan::flat_update: blockops::cpu::flat_update(lscale, lhs, rscale, rhs); break; - case Plan::reorder: - blockops::reorder::index_permute_reorder(lscale, lhs.buf(), rscale, rhs.buf(), + case Plan::permute: + blockops::permute::index_permute(lscale, lhs.buf(), rscale, rhs.buf(), ip_plan_.perm_, rhs.block_dims()); break; case Plan::index_permute_assign: NOT_ALLOWED(); break; @@ -265,12 +265,12 @@ class BlockAssignPlan { } template - void prep_reorder(const std::vector& lhs_labels, const std::vector& rhs_labels) { + void prep_permute(const std::vector& lhs_labels, const std::vector& rhs_labels) { prep_index_permute_plan(lhs_labels, rhs_labels); if(plan_ == Plan::invalid) { - return; // no index permute plan, no reorder plan + return; // no index permute plan, no permute plan } - plan_ = Plan::reorder; + plan_ = Plan::permute; } template @@ -315,7 +315,7 @@ class BlockAssignPlan { enum class Plan { flat_assign, flat_update, - reorder, + permute, index_permute_assign, index_permute_update, ipgen_loop_assign, @@ -343,7 +343,7 @@ class BlockAssignPlan { switch(plan) { case Plan::flat_assign: return "flat_assign"; case Plan::flat_update: return "flat_update"; - case Plan::reorder: return "reorder"; + case Plan::permute: return "permute"; case Plan::index_permute_assign: return "index_permute_assign"; case Plan::index_permute_update: return "index_permute_update"; case Plan::ipgen_loop_assign: return "ipgen_loop_assign"; diff --git a/src/tamm/ip_reorder.hpp b/src/tamm/ip_permute.hpp similarity index 59% rename from src/tamm/ip_reorder.hpp rename to src/tamm/ip_permute.hpp index caeb14143..926da53e0 100644 --- a/src/tamm/ip_reorder.hpp +++ b/src/tamm/ip_permute.hpp @@ -4,42 +4,41 @@ #include #include "tamm/errors.hpp" -#include "tamm/kernels/cpu_reorder.hpp" +#include "tamm/kernels/cpu_permute.hpp" #include "tamm/types.hpp" -namespace tamm::blockops::reorder { +namespace tamm::blockops::permute { /////////////////////////////////////////////////////////////////////////////// // -// Index permute using the in-house CPU reorder kernel +// Index permute using the in-house CPU permute kernel // (replaces the old HPTT-based index_permute_hptt) // /////////////////////////////////////////////////////////////////////////////// -// Out-of-place transpose with HPTT-identical semantics: -// lbuf = lscale * lbuf + rscale * transpose(rbuf) +// Out-of-place permute with HPTT-identical semantics: +// lbuf = lscale * lbuf + rscale * permute(rbuf) // perm_to_dest maps each destination axis to its source axis and sdims holds // the source extents, both in natural order; the output extents follow as // outDims[k] == sdims[perm_to_dest[k]]. Scale types are independent template // parameters so call sites may mix e.g. an int literal rscale with a // complex buffer (as the old single-T HPTT wrapper allowed via conversion). template -void index_permute_reorder(BL lscale, T1* lbuf, BA rscale, const T2* rbuf, - const PermVector& perm_to_dest, const std::vector& sdims) { +void index_permute(BL lscale, T1* lbuf, BA rscale, const T2* rbuf, + const PermVector& perm_to_dest, const std::vector& sdims) { const size_t ndim = sdims.size(); EXPECTS(perm_to_dest.size() == ndim); - EXPECTS(ndim <= static_cast(kernels::gpu::reorder_maxrank)); + EXPECTS(ndim <= static_cast(kernels::gpu::permute_maxrank)); EXPECTS(lbuf != nullptr && (ndim == 0 || rbuf != nullptr)); - size_t outDims[kernels::gpu::reorder_maxrank] = {}; - int perm[kernels::gpu::reorder_maxrank] = {}; + size_t outDims[kernels::gpu::permute_maxrank] = {}; + int perm[kernels::gpu::permute_maxrank] = {}; for(size_t i = 0; i < ndim; ++i) { perm[i] = static_cast(perm_to_dest[i]); EXPECTS(perm[i] >= 0 && static_cast(perm[i]) < ndim); outDims[i] = sdims[static_cast(perm[i])]; } - kernels::cpu::transpose_reorder_cpu(lbuf, rbuf, static_cast(ndim), outDims, perm, rscale, - lscale); + kernels::cpu::permute(lbuf, rbuf, static_cast(ndim), outDims, perm, rscale, lscale); } -} // namespace tamm::blockops::reorder +} // namespace tamm::blockops::permute diff --git a/src/tamm/kernels/assign.hpp b/src/tamm/kernels/assign.hpp index c34dea08d..1c870c429 100644 --- a/src/tamm/kernels/assign.hpp +++ b/src/tamm/kernels/assign.hpp @@ -1,7 +1,7 @@ #pragma once #include "tamm/errors.hpp" -#include "tamm/kernels/cpu_reorder.hpp" +#include "tamm/kernels/cpu_permute.hpp" #include "tamm/types.hpp" #include "tamm/utils.hpp" @@ -373,16 +373,16 @@ void ip_gen_loop(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T sca } template -void ip_reorder(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, +void ip_permute(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign = true) { const size_t ndim = ddims.size(); EXPECTS(sdims.size() == ndim && dlabels.size() == ndim && slabels.size() == ndim); - EXPECTS(ndim <= static_cast(kernels::gpu::reorder_maxrank)); + EXPECTS(ndim <= static_cast(kernels::gpu::permute_maxrank)); // Natural-order output extents + output-axis -> source-axis map, exactly // the (size, perm) pair the old HPTT plan consumed. - size_t outDims[kernels::gpu::reorder_maxrank] = {}; - int perm[kernels::gpu::reorder_maxrank] = {}; + size_t outDims[kernels::gpu::permute_maxrank] = {}; + int perm[kernels::gpu::permute_maxrank] = {}; for(size_t i = 0; i < ndim; i++) { auto it = std::find(slabels.begin(), slabels.end(), dlabels[i]); EXPECTS(it != slabels.end()); @@ -390,7 +390,7 @@ void ip_reorder(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scal perm[i] = static_cast(j); outDims[i] = sdims[j].value(); } - kernels::cpu::transpose_reorder_cpu(dst, src, static_cast(ndim), outDims, perm, scale, + kernels::cpu::permute(dst, src, static_cast(ndim), outDims, perm, scale, is_assign ? T{0} : T{1}); } @@ -473,7 +473,7 @@ void assign(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, c if(is_assign) { internal::index_permute(dst, src, perm_to_dest, ddims, scale); } else { internal::index_permute_acc(dst, src, perm_to_dest, ddims, scale); } } - else internal::ip_reorder(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); + else internal::ip_permute(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } else { internal::ip_gen_loop(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } } diff --git a/src/tamm/kernels/cpu_reorder.hpp b/src/tamm/kernels/cpu_permute.hpp similarity index 56% rename from src/tamm/kernels/cpu_reorder.hpp rename to src/tamm/kernels/cpu_permute.hpp index 14e02a1cf..eada57b0e 100644 --- a/src/tamm/kernels/cpu_reorder.hpp +++ b/src/tamm/kernels/cpu_permute.hpp @@ -1,18 +1,18 @@ #pragma once -// Host-only CPU tensor transpose, the HPTT replacement. +// Host-only CPU tensor permute, the HPTT replacement. // -// Row-major (TAMM-native) counterpart of the GPU kernel in gpu_reorder.hpp. -// Reuses its ReorderMeta storage, elementwise scale/add helpers and identity +// Row-major (TAMM-native) counterpart of the GPU kernel in gpu_permute.hpp. +// Reuses its PermuteMeta storage, elementwise scale/add helpers and identity // check, so both backends implement the same math; only the metadata build // differs (row-major strides, natural axis order). // -// out = beta * out + alpha * transpose(in) +// out = beta * out + alpha * permute(in) // // beta==0 never reads out; beta==1 accumulates. Single-threaded and // allocation-free except on the aliased in-place path. -#include "tamm/kernels/gpu_reorder.hpp" +#include "tamm/kernels/gpu_permute.hpp" #include #include @@ -23,8 +23,8 @@ namespace tamm::kernels::cpu { // Split a host-side alpha/beta into raw {real, imag} doubles. template -inline void reorder_split_scale_host(S scale, double& re, double& im) { - if constexpr(gpu::reorder_is_complex_v) { +inline void permute_split_scale_host(S scale, double& re, double& im) { + if constexpr(gpu::permute_is_complex_v) { re = static_cast(scale.real()); im = static_cast(scale.imag()); } @@ -35,27 +35,27 @@ inline void reorder_split_scale_host(S scale, double& re, double& im) { } template -inline bool reorder_scale_is_zero_host(S scale) { - if constexpr(gpu::reorder_is_complex_v) { return scale.real() == 0 && scale.imag() == 0; } +inline bool permute_scale_is_zero_host(S scale) { + if constexpr(gpu::permute_is_complex_v) { return scale.real() == 0 && scale.imag() == 0; } else { return scale == S{0}; } } template -inline bool reorder_scale_is_one_host(S scale) { - if constexpr(gpu::reorder_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } +inline bool permute_scale_is_one_host(S scale) { + if constexpr(gpu::permute_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } else { return scale == S{1}; } } // Build row-major metadata from natural output extents + dst->src perm. // Source extents follow from inDims[perm[k]] == outDims[k]. -inline gpu::ReorderMeta reorder_build_meta_rowmajor(int ndim, const size_t* outDims, +inline gpu::PermuteMeta permute_build_meta_rowmajor(int ndim, const size_t* outDims, const int* perm) { - EXPECTS(ndim >= 0 && ndim <= gpu::reorder_maxrank); - gpu::ReorderMeta meta{}; + EXPECTS(ndim >= 0 && ndim <= gpu::permute_maxrank); + gpu::PermuteMeta meta{}; meta.ndim = ndim; if(ndim == 0) return meta; - size_t inDims[gpu::reorder_maxrank] = {}; + size_t inDims[gpu::permute_maxrank] = {}; for(int k = 0; k < ndim; ++k) { EXPECTS(perm[k] >= 0 && perm[k] < ndim); meta.outDims[k] = outDims[k]; @@ -72,24 +72,24 @@ inline gpu::ReorderMeta reorder_build_meta_rowmajor(int ndim, const size_t* outD } template -inline void transpose_reorder_cpu_impl(T* out, const T* in, const gpu::ReorderMeta& meta, +inline void permute_impl(T* out, const T* in, const gpu::PermuteMeta& meta, size_t total, double alpha_re, double alpha_im, B beta) { - const bool beta_is_zero = reorder_scale_is_zero_host(beta); - const bool beta_is_one = reorder_scale_is_one_host(beta); + const bool beta_is_zero = permute_scale_is_zero_host(beta); + const bool beta_is_one = permute_scale_is_one_host(beta); double beta_re = 0.0, beta_im = 0.0; - if(!beta_is_zero && !beta_is_one) reorder_split_scale_host(beta, beta_re, beta_im); + if(!beta_is_zero && !beta_is_one) permute_split_scale_host(beta, beta_re, beta_im); // Odometer over the output coordinates: no division/modulo per element, // just ndim multiply-adds plus an amortized single increment. const int ndim = meta.ndim; - size_t coord[gpu::reorder_maxrank] = {}; + size_t coord[gpu::permute_maxrank] = {}; for(size_t t = 0; t < total; ++t) { size_t src = 0; for(int k = 0; k < ndim; ++k) src += coord[k] * meta.inStrides[meta.perm[k]]; - const T y = gpu::reorder_scaled(in[src], alpha_re, alpha_im); + const T y = gpu::permute_scaled(in[src], alpha_re, alpha_im); if(beta_is_zero) { out[t] = y; } - else if(beta_is_one) { out[t] = gpu::reorder_add(out[t], y); } - else { out[t] = gpu::reorder_add(gpu::reorder_scaled(out[t], beta_re, beta_im), y); } + else if(beta_is_one) { out[t] = gpu::permute_add(out[t], y); } + else { out[t] = gpu::permute_add(gpu::permute_scaled(out[t], beta_re, beta_im), y); } for(int k = ndim - 1; k >= 0; --k) { if(++coord[k] < meta.outDims[k]) break; coord[k] = 0; @@ -97,44 +97,44 @@ inline void transpose_reorder_cpu_impl(T* out, const T* in, const gpu::ReorderMe } } -// out = beta * out + alpha * transpose(in); outDims/perm in natural order. +// out = beta * out + alpha * permute(in); outDims/perm in natural order. template -inline void transpose_reorder_cpu(T* out, const T* in, int ndim, const size_t* outDims, +inline void permute(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, A alpha, B beta) { - EXPECTS(ndim >= 0 && ndim <= gpu::reorder_maxrank); + EXPECTS(ndim >= 0 && ndim <= gpu::permute_maxrank); if(ndim == 0) { double alpha_re, alpha_im, beta_re, beta_im; - reorder_split_scale_host(alpha, alpha_re, alpha_im); - reorder_split_scale_host(beta, beta_re, beta_im); - const T y = gpu::reorder_scaled(in[0], alpha_re, alpha_im); - if(reorder_scale_is_zero_host(beta)) { out[0] = y; } - else { out[0] = gpu::reorder_add(gpu::reorder_scaled(out[0], beta_re, beta_im), y); } + permute_split_scale_host(alpha, alpha_re, alpha_im); + permute_split_scale_host(beta, beta_re, beta_im); + const T y = gpu::permute_scaled(in[0], alpha_re, alpha_im); + if(permute_scale_is_zero_host(beta)) { out[0] = y; } + else { out[0] = gpu::permute_add(gpu::permute_scaled(out[0], beta_re, beta_im), y); } return; } EXPECTS(outDims != nullptr && perm != nullptr && in != nullptr && out != nullptr); - const gpu::ReorderMeta meta = reorder_build_meta_rowmajor(ndim, outDims, perm); - const size_t total = gpu::reorder_total(ndim, outDims); + const gpu::PermuteMeta meta = permute_build_meta_rowmajor(ndim, outDims, perm); + const size_t total = gpu::permute_total(ndim, outDims); if(total == 0) return; double alpha_re, alpha_im; - reorder_split_scale_host(alpha, alpha_re, alpha_im); + permute_split_scale_host(alpha, alpha_re, alpha_im); // Identity + alpha==1 + beta==0: pure copy. - if(reorder_scale_is_zero_host(beta) && reorder_scale_is_one_host(alpha) && - gpu::reorder_is_identity(meta)) { + if(permute_scale_is_zero_host(beta) && permute_scale_is_one_host(alpha) && + gpu::permute_is_identity(meta)) { std::memcpy(out, in, total * sizeof(T)); return; } // Aliased non-identity permute goes via a temp buffer. - if(out == in && !gpu::reorder_is_identity(meta)) { + if(out == in && !gpu::permute_is_identity(meta)) { std::vector tmp(in, in + total); - transpose_reorder_cpu_impl(out, tmp.data(), meta, total, alpha_re, alpha_im, beta); + permute_impl(out, tmp.data(), meta, total, alpha_re, alpha_im, beta); return; } - transpose_reorder_cpu_impl(out, in, meta, total, alpha_re, alpha_im, beta); + permute_impl(out, in, meta, total, alpha_re, alpha_im, beta); } } // namespace tamm::kernels::cpu diff --git a/src/tamm/kernels/gpu_reorder.cpp b/src/tamm/kernels/gpu_permute.cpp similarity index 65% rename from src/tamm/kernels/gpu_reorder.cpp rename to src/tamm/kernels/gpu_permute.cpp index 331e9e3e7..d95be018a 100644 --- a/src/tamm/kernels/gpu_reorder.cpp +++ b/src/tamm/kernels/gpu_permute.cpp @@ -1,4 +1,4 @@ -#include "gpu_reorder.hpp" +#include "gpu_permute.hpp" #include "tamm_blas.hpp" @@ -9,8 +9,8 @@ namespace tamm::kernels::gpu { // Decompose the host-side scale into raw {real, imag} doubles for the kernel. template -inline void reorder_split_scale(T scale, double& re, double& im) { - if constexpr(reorder_is_complex_v) { +inline void permute_split_scale(T scale, double& re, double& im) { + if constexpr(permute_is_complex_v) { re = static_cast(scale.real()); im = static_cast(scale.imag()); } @@ -21,8 +21,8 @@ inline void reorder_split_scale(T scale, double& re, double& im) { } template -inline bool reorder_scale_is_one(T scale) { - if constexpr(reorder_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } +inline bool permute_scale_is_one(T scale) { + if constexpr(permute_is_complex_v) { return scale.real() == 1 && scale.imag() == 0; } else { return scale == T{1}; } } @@ -35,60 +35,60 @@ __global__ void #if defined(USE_CUDA) || defined(USE_HIP) __launch_bounds__(256) #endif -reorder_kernel(T* out, const T* in, ReorderMeta meta, size_t total, double scale_re, double scale_im, +permute_kernel(T* out, const T* in, PermuteMeta meta, size_t total, double scale_re, double scale_im, bool accumulate) { const size_t tid0 = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const size_t stride = static_cast(gridDim.x) * blockDim.x; for(size_t t = tid0; t < total; t += stride) { const Idx tid = static_cast(t); - const size_t src = reorder_src_index(tid, meta); - const T y = reorder_scaled(in[src], scale_re, scale_im); - out[tid] = accumulate ? reorder_add(out[tid], y) : y; + const size_t src = permute_src_index(tid, meta); + const T y = permute_scaled(in[src], scale_re, scale_im); + out[tid] = accumulate ? permute_add(out[tid], y) : y; } } #endif // USE_CUDA || USE_HIP template -void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, +void permute(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, T scale, bool accumulate, gpuStream_t& handle) { - EXPECTS(ndim >= 0 && ndim <= reorder_maxrank); + EXPECTS(ndim >= 0 && ndim <= permute_maxrank); EXPECTS(out != in); // out-of-place only; in-place permutes need a temp buffer if(ndim == 0) { // Rank-0: single element, still honors scale/accumulate. - if(reorder_scale_is_one(scale) && !accumulate) { + if(permute_scale_is_one(scale) && !accumulate) { gpuMemcpyAsync(out, in, 1, gpuMemcpyDeviceToDevice, handle); return; } double scale_re, scale_im; - reorder_split_scale(scale, scale_re, scale_im); + permute_split_scale(scale, scale_re, scale_im); #if defined(USE_CUDA) || defined(USE_HIP) // Route through the kernel for stream ordering (async, like every // other path through this function). - ReorderMeta meta{}; + PermuteMeta meta{}; meta.ndim = 0; - reorder_kernel + permute_kernel <<<1, 1, 0, handle.first>>>(out, in, meta, 1, scale_re, scale_im, accumulate); #elif defined(USE_DPCPP) handle.first.parallel_for(sycl::range<1>(1), [=](sycl::id<1>) { - const T y = reorder_scaled(in[0], scale_re, scale_im); - out[0] = accumulate ? reorder_add(out[0], y) : y; + const T y = permute_scaled(in[0], scale_re, scale_im); + out[0] = accumulate ? permute_add(out[0], y) : y; }); #endif return; } EXPECTS(outDims != nullptr && perm != nullptr); - const ReorderMeta meta = reorder_build_meta(ndim, outDims, perm); - const size_t total = reorder_total(ndim, outDims); + const PermuteMeta meta = permute_build_meta(ndim, outDims, perm); + const size_t total = permute_total(ndim, outDims); if(total == 0) return; double scale_re, scale_im; - reorder_split_scale(scale, scale_re, scale_im); - if constexpr(!reorder_is_complex_v) { EXPECTS(scale_im == 0.0); } + permute_split_scale(scale, scale_re, scale_im); + if constexpr(!permute_is_complex_v) { EXPECTS(scale_im == 0.0); } // Identity permutation with scale==1 and overwrite: pure copy, no kernel. - if(!accumulate && reorder_scale_is_one(scale) && reorder_is_identity(meta)) { + if(!accumulate && permute_scale_is_one(scale) && permute_is_identity(meta)) { gpuMemcpyAsync(out, in, total, gpuMemcpyDeviceToDevice, handle); return; } @@ -99,31 +99,31 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con constexpr size_t block = 256; constexpr size_t maxBlocks = 1 << 20; // stride loop covers the rest const size_t nblocks = std::min((total + block - 1) / block, maxBlocks); - if(reorder_meta_fits32(meta, total)) { - reorder_kernel<<(nblocks), static_cast(block), 0, + if(permute_meta_fits32(meta, total)) { + permute_kernel<<(nblocks), static_cast(block), 0, handle.first>>>(out, in, meta, total, scale_re, scale_im, accumulate); } else { - reorder_kernel<<(nblocks), static_cast(block), 0, + permute_kernel<<(nblocks), static_cast(block), 0, handle.first>>>(out, in, meta, total, scale_re, scale_im, accumulate); } #elif defined(USE_DPCPP) - if(reorder_meta_fits32(meta, total)) { + if(permute_meta_fits32(meta, total)) { handle.first.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { const uint32_t tid = static_cast(idx[0]); - const size_t src = reorder_src_index(tid, meta); - const T y = reorder_scaled(in[src], scale_re, scale_im); - out[tid] = accumulate ? reorder_add(out[tid], y) : y; + const size_t src = permute_src_index(tid, meta); + const T y = permute_scaled(in[src], scale_re, scale_im); + out[tid] = accumulate ? permute_add(out[tid], y) : y; }); } else { handle.first.parallel_for(sycl::range<1>(total), [=](sycl::id<1> idx) { const size_t t = idx[0]; - const size_t src = reorder_src_index(t, meta); - const T y = reorder_scaled(in[src], scale_re, scale_im); - out[t] = accumulate ? reorder_add(out[t], y) : y; + const size_t src = permute_src_index(t, meta); + const T y = permute_scaled(in[src], scale_re, scale_im); + out[t] = accumulate ? permute_add(out[t], y) : y; }); } // No wait: the caller's in-order queue preserves ordering, and @@ -131,16 +131,16 @@ void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, con #endif } -template void transpose_reorder(double* out, const double* in, int ndim, const size_t* outDims, +template void permute(double* out, const double* in, int ndim, const size_t* outDims, const int* perm, double scale, bool accumulate, gpuStream_t& handle); -template void transpose_reorder(float* out, const float* in, int ndim, const size_t* outDims, +template void permute(float* out, const float* in, int ndim, const size_t* outDims, const int* perm, float scale, bool accumulate, gpuStream_t& handle); -template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, +template void permute(std::complex* out, const std::complex* in, int ndim, const size_t* outDims, const int* perm, std::complex scale, bool accumulate, gpuStream_t& handle); -template void transpose_reorder(std::complex* out, const std::complex* in, int ndim, +template void permute(std::complex* out, const std::complex* in, int ndim, const size_t* outDims, const int* perm, std::complex scale, bool accumulate, gpuStream_t& handle); diff --git a/src/tamm/kernels/gpu_reorder.hpp b/src/tamm/kernels/gpu_permute.hpp similarity index 68% rename from src/tamm/kernels/gpu_reorder.hpp rename to src/tamm/kernels/gpu_permute.hpp index b9148e881..6071b3b5a 100644 --- a/src/tamm/kernels/gpu_reorder.hpp +++ b/src/tamm/kernels/gpu_permute.hpp @@ -1,24 +1,24 @@ #pragma once -// Host/device shared logic for TAMM's in-house GPU tensor transpose (reorder). +// Host/device shared logic for TAMM's in-house GPU tensor permute. // -// This header holds everything about the transpose that is backend-agnostic: +// This header holds everything about the permute that is backend-agnostic: // - the permutation metadata layout, // - the output-linear-index -> input-linear-index decode (column-major, // first-axis-fastest, matching the convention the old librett path used), // - the scale/accumulate elementwise apply, // - host-side helpers that build the metadata from TAMM's (dims, labels). // -// The per-backend kernels in gpu_reorder.cpp are thin launch wrappers around +// The per-backend kernels in gpu_permute.cpp are thin launch wrappers around // these helpers, so all three backends (CUDA/HIP/SYCL) provably implement the -// same math. The CPU unit test (tests/tamm/Test_TransposeReorder.cpp) +// same math. The CPU unit test (tests/tamm/Test_TransposePermute.cpp) // exercises these helpers directly, without needing a GPU. // // Layout note: TAMM callers pass dims/labels in reverse (Fortran) order -// (see build_reorder_spec); the kernel itself always works in column-major +// (see build_permute_spec); the kernel itself always works in column-major // (first-axis-fastest) order. Reverse + column-major is exactly equivalent to -// the natural row-major (last-axis-fastest) transpose, which is what the CPU -// reference (kernels::assign / HPTT path) computes. +// the natural row-major (last-axis-fastest) permute, which is what the CPU +// reference (kernels::assign path) computes. #include #include @@ -32,48 +32,48 @@ #include "tamm/types.hpp" #if defined(__CUDACC__) || defined(__HIPCC__) -#define TAMM_REORDER_HD __host__ __device__ +#define TAMM_PERMUTE_HD __host__ __device__ #else -#define TAMM_REORDER_HD +#define TAMM_PERMUTE_HD #endif namespace tamm::kernels::gpu { -// Maximum tensor rank supported by the reorder kernel. -inline constexpr int reorder_maxrank = 8; -static_assert(reorder_maxrank >= static_cast(tamm::maxrank), - "reorder kernel must support at least tamm::maxrank"); +// Maximum tensor rank supported by the permute kernel. +inline constexpr int permute_maxrank = 8; +static_assert(permute_maxrank >= static_cast(tamm::maxrank), + "permute kernel must support at least tamm::maxrank"); template -struct reorder_is_complex: std::false_type {}; +struct permute_is_complex: std::false_type {}; template -struct reorder_is_complex>: std::true_type {}; +struct permute_is_complex>: std::true_type {}; template -inline constexpr bool reorder_is_complex_v = reorder_is_complex::value; +inline constexpr bool permute_is_complex_v = permute_is_complex::value; // POD metadata passed by value into the device kernel. Maps an output linear // index to the corresponding source linear index for an arbitrary axis // permutation. All extents/strides are 64-bit: TAMM blocks can exceed 2^31 // elements (cf. the B*K*N = 2^32 overflow previously fixed in gemm_wrapper). -struct ReorderMeta { +struct PermuteMeta { int ndim = 0; - size_t outDims[reorder_maxrank] = {}; - size_t outStrides[reorder_maxrank] = {}; - size_t inStrides[reorder_maxrank] = {}; - int perm[reorder_maxrank] = {}; // output-axis -> source-axis map + size_t outDims[permute_maxrank] = {}; + size_t outStrides[permute_maxrank] = {}; + size_t inStrides[permute_maxrank] = {}; + int perm[permute_maxrank] = {}; // output-axis -> source-axis map }; // Passed by value into CUDA/HIP kernels (4 KB parameter limit) and captured // by value into SYCL kernels (must stay trivially copyable). -static_assert(sizeof(ReorderMeta) <= 1024, "ReorderMeta must stay small"); -static_assert(std::is_trivially_copyable_v, - "ReorderMeta must be trivially copyable for device kernels"); +static_assert(sizeof(PermuteMeta) <= 1024, "PermuteMeta must stay small"); +static_assert(std::is_trivially_copyable_v, + "PermuteMeta must be trivially copyable for device kernels"); // Decode: output linear index -> source linear index (column-major). // Templated on the integer type so launches with < 2^32 elements can use // 32-bit division/modulo (much faster on GPUs); the caller guarantees every // value used below fits in Idx (see meta_fits32). template -TAMM_REORDER_HD inline size_t reorder_src_index(Idx tid, const ReorderMeta& meta) { +TAMM_PERMUTE_HD inline size_t permute_src_index(Idx tid, const PermuteMeta& meta) { size_t src = 0; for(int k = 0; k < meta.ndim; ++k) { const Idx oidx = @@ -91,8 +91,8 @@ TAMM_REORDER_HD inline size_t reorder_src_index(Idx tid, const ReorderMeta& meta // device calls are needed; the complex <-> T[2] pun is blessed by the // standard. template -TAMM_REORDER_HD inline T reorder_scaled(T v, double scale_re, double scale_im) { - if constexpr(reorder_is_complex_v) { +TAMM_PERMUTE_HD inline T permute_scaled(T v, double scale_re, double scale_im) { + if constexpr(permute_is_complex_v) { using R = typename T::value_type; const R* w = reinterpret_cast(&v); T out; @@ -107,8 +107,8 @@ TAMM_REORDER_HD inline T reorder_scaled(T v, double scale_re, double scale_im) { } template -TAMM_REORDER_HD inline T reorder_add(T a, T b) { - if constexpr(reorder_is_complex_v) { +TAMM_PERMUTE_HD inline T permute_add(T a, T b) { + if constexpr(permute_is_complex_v) { using R = typename T::value_type; const R* wa = reinterpret_cast(&a); const R* wb = reinterpret_cast(&b); @@ -124,13 +124,13 @@ TAMM_REORDER_HD inline T reorder_add(T a, T b) { } // Per-element store idiom (use at every call site): -// const T y = reorder_scaled(in[src], scale_re, scale_im); -// out[tid] = accumulate ? reorder_add(out[tid], y) : y; +// const T y = permute_scaled(in[src], scale_re, scale_im); +// out[tid] = accumulate ? permute_add(out[tid], y) : y; // so the overwrite path never reads the output buffer and no code path // default-constructs a T on the device. // True when every extent/stride/total used by the decode fits in 32 bits. -inline bool reorder_meta_fits32(const ReorderMeta& meta, size_t total) { +inline bool permute_meta_fits32(const PermuteMeta& meta, size_t total) { constexpr auto max32 = static_cast(std::numeric_limits::max()); if(total > max32) return false; for(int k = 0; k < meta.ndim; ++k) { @@ -141,7 +141,7 @@ inline bool reorder_meta_fits32(const ReorderMeta& meta, size_t total) { return true; } -inline bool reorder_is_identity(const ReorderMeta& meta) { +inline bool permute_is_identity(const PermuteMeta& meta) { for(int k = 0; k < meta.ndim; ++k) { if(meta.perm[k] != k) return false; } @@ -151,13 +151,13 @@ inline bool reorder_is_identity(const ReorderMeta& meta) { // Host-side: build the metadata from the (already reversed) output extents // and output-axis -> source-axis permutation. Source extents follow from // inDims[perm[k]] == outDims[k]; both tensors are contiguous column-major. -inline ReorderMeta reorder_build_meta(int ndim, const size_t* outDims, const int* perm) { - EXPECTS(ndim >= 0 && ndim <= reorder_maxrank); - ReorderMeta meta{}; +inline PermuteMeta permute_build_meta(int ndim, const size_t* outDims, const int* perm) { + EXPECTS(ndim >= 0 && ndim <= permute_maxrank); + PermuteMeta meta{}; meta.ndim = ndim; if(ndim == 0) return meta; - size_t inDims[reorder_maxrank] = {}; + size_t inDims[permute_maxrank] = {}; for(int k = 0; k < ndim; ++k) { EXPECTS(perm[k] >= 0 && perm[k] < ndim); meta.outDims[k] = outDims[k]; @@ -173,7 +173,7 @@ inline ReorderMeta reorder_build_meta(int ndim, const size_t* outDims, const int return meta; } -inline size_t reorder_total(int ndim, const size_t* outDims) { +inline size_t permute_total(int ndim, const size_t* outDims) { size_t total = 1; for(int k = 0; k < ndim; ++k) total *= outDims[k]; return total; @@ -182,28 +182,28 @@ inline size_t reorder_total(int ndim, const size_t* outDims) { // Host-side: translate TAMM's (src dims, src labels, dst labels) into the // reversed (Fortran-order) output extents + permutation the column-major // kernel consumes. Reversing both dim orders converts between TAMM's natural -// row-major transpose and the kernel's column-major transpose. +// row-major permute and the kernel's column-major permute. // // Allocation-free: everything lives in fixed-size stack arrays (rank <= -// reorder_maxrank), so this can sit in the hottest transpose path without +// permute_maxrank), so this can sit in the hottest permute path without // touching the heap. (An earlier version used std::vectors here: 4 heap -// allocations per transpose call.) -struct ReorderSpec { +// allocations per permute call.) +struct PermuteSpec { int ndim = 0; - size_t outDims[reorder_maxrank] = {}; - int perm[reorder_maxrank] = {}; // reversed output-axis -> source-axis map + size_t outDims[permute_maxrank] = {}; + int perm[permute_maxrank] = {}; // reversed output-axis -> source-axis map }; -inline void build_reorder_spec(const SizeVec& sdims, const IntLabelVec& slabels, - const IntLabelVec& dlabels, ReorderSpec& spec) { +inline void build_permute_spec(const SizeVec& sdims, const IntLabelVec& slabels, + const IntLabelVec& dlabels, PermuteSpec& spec) { const size_t ndim = sdims.size(); EXPECTS(dlabels.size() == ndim && slabels.size() == ndim); - EXPECTS(ndim <= static_cast(reorder_maxrank)); + EXPECTS(ndim <= static_cast(permute_maxrank)); // Reverse into Fortran order (librett used the same convention). - size_t r_size[reorder_maxrank]; - IntLabel r_slabels[reorder_maxrank]; - IntLabel r_dlabels[reorder_maxrank]; + size_t r_size[permute_maxrank]; + IntLabel r_slabels[permute_maxrank]; + IntLabel r_dlabels[permute_maxrank]; for(size_t i = 0; i < ndim; ++i) { r_size[i] = sdims[ndim - 1 - i].value(); r_slabels[i] = slabels[ndim - 1 - i]; diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index e8e91f64f..eb2e29eaa 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -16,7 +16,7 @@ #include "tamm/rmm_memory_pool.hpp" #include "tamm/utils.hpp" #include "tamm_blas.hpp" -#include "gpu_reorder.hpp" +#include "gpu_permute.hpp" #if !defined(USE_CUDA) && !defined(USE_HIP) && !defined(USE_DPCPP) namespace tamm { @@ -28,11 +28,11 @@ namespace tamm { namespace kernels { -// Device staging owned by one transpose_inputs() call. Type-erased to bytes +// Device staging owned by one permute_inputs() call. Type-erased to bytes // (the pool frees exactly size_bytes()) so block_multiply() can hold it in a // single function-scope variable across its differently-typed branches and // return it to the pool under the end-of-call drain. -struct GpuTransposeStaging { +struct GpuPermuteStaging { std::span a{}; std::span b{}; }; @@ -156,26 +156,26 @@ void assign_gpu(gpuStream_t& thandle, T*& dst, const SizeVec& ddims, const IntLa const Size ssize = std::accumulate(sdims.begin(), sdims.end(), Size{1}, std::multiplies()); if(ssize.value() == 0) return; - EXPECTS(ndim <= static_cast(gpu::reorder_maxrank)); + EXPECTS(ndim <= static_cast(gpu::permute_maxrank)); - // Fortran-order (reversed) spec for the column-major reorder kernel; this - // is exactly the row-major transpose the CPU assign() computes. + // Fortran-order (reversed) spec for the column-major permute kernel; this + // is exactly the row-major permute the CPU assign() computes. // Allocation-free (stack spec): no heap traffic in this hot path. - gpu::ReorderSpec spec{}; - gpu::build_reorder_spec(sdims, slabels, dlabels, spec); - gpu::transpose_reorder(dst, src, static_cast(ndim), spec.outDims, spec.perm, scale, + gpu::PermuteSpec spec{}; + gpu::build_permute_spec(sdims, slabels, dlabels, spec); + gpu::permute(dst, src, static_cast(ndim), spec.outDims, spec.perm, scale, !is_assign, thandle); } #endif template -bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, +bool permute_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, const SizeVec& ainter_dims, const IntLabelVec& ainter_labels, const T2* abuf, size_t asize, const SizeVec& adims, const IntLabelVec& alabels, T3* binter_buf, const SizeVec& binter_dims, const IntLabelVec& binter_labels, const T3* bbuf, size_t bsize, const SizeVec& bdims, const IntLabelVec& blabels, T2*& ainter_buf_dev, T3*& binter_buf_dev, - GpuTransposeStaging& staging) { + GpuPermuteStaging& staging) { bool gpu_trans = false; #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) @@ -193,7 +193,7 @@ bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, assign_gpu(thandle, binter_buf_dev, binter_dims, binter_labels, T3{1}, binter_dev_in.data(), bdims, blabels, true); - // Staging is NOT freed here: it stays alive across GEMM/output-transpose + // Staging is NOT freed here: it stays alive across GEMM/output-permute // and is returned to the pool just before block_multiply()'s single end // drain, saving one device-wide sync per call. (On the in-order queue all // of the above is ordered behind prior work by construction.) @@ -210,7 +210,7 @@ bool transpose_inputs(ExecutionHW hw, gpuStream_t& thandle, T2* ainter_buf, } template -void transpose_output(ExecutionHW hw, gpuStream_t& thandle, bool gpu_trans, T1* cinter_buf, +void permute_output(ExecutionHW hw, gpuStream_t& thandle, bool gpu_trans, T1* cinter_buf, const SizeVec& cinter_dims, const IntLabelVec& cinter_labels, T1* cbuf, const SizeVec& cdims, const IntLabelVec& clabels, T1*& cinter_buf_dev, T1*& cinter_tmp_buf_dev, bool is_assign) { @@ -220,11 +220,11 @@ void transpose_output(ExecutionHW hw, gpuStream_t& thandle, bool gpu_trans, T1* // old librett assign_gpu semantics (which ignored is_assign/scale). // Accumulation across reduction iterations here is performed by GEMM // itself (beta=cscale=1 accumulates in cinter_tmp_buf_dev in inter - // ordering), so the output transpose must overwrite cinter_buf_dev with + // ordering), so the output permute must overwrite cinter_buf_dev with // the running total each iteration. Honoring is_assign=false here (D+=T) // double-counts: iter2 adds T1+S2 on top of D=T1, giving 2*S1+S2. // (CPU path differs: its per-call cinter_buf is freshly zeroed, so the - // transpose is the accumulator there and must honor is_assign.) + // permute is the accumulator there and must honor is_assign.) assign_gpu(thandle, cinter_buf_dev, cdims, clabels, T1{1}, cinter_tmp_buf_dev, cinter_dims, cinter_labels, true); return; @@ -365,8 +365,8 @@ void block_multiply( bool gpu_trans = false; // Input staging for the GPU path; returned to the pool under the - // end-of-call drain (see transpose_inputs). - GpuTransposeStaging transpose_staging{}; + // end-of-call drain (see permute_inputs). + GpuPermuteStaging permute_staging{}; std::span cinter_span = allocate_host_buffer(hw, static_cast(csize.value())); T1* cinter_buf = cinter_span.data(); @@ -389,10 +389,10 @@ void block_multiply( T2* ainter_buf = ainter_span.data(); T3* binter_buf = binter_span.data(); - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, ainter_buf_dev, - binter_buf_dev, transpose_staging); + binter_buf_dev, permute_staging); if(!gpu_trans) copy_data_to_gpu(hw, thandle, ainter_buf, asize.value(), ainter_buf_dev, binter_buf, @@ -401,7 +401,7 @@ void block_multiply( gemm_wrapper(hw, thandle, AR, BR, B, M, N, K, alpha, beta, ainter_buf, ainter_buf_dev, binter_buf, binter_buf_dev, cinter_buf, cinter_tmp_buf_dev); - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, cinter_tmp_buf_dev, is_assign); free_host_buffer(hw, ainter_span); @@ -430,10 +430,10 @@ void block_multiply( std::span bbuf_complex_dev_span = allocate_device_buffer(hw, bsize.value()); T1* bbuf_complex_dev = bbuf_complex_dev_span.data(); - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf_complex, bsize.value(), bdims, blabels, - ainter_buf_dev, bbuf_complex_dev, transpose_staging); + ainter_buf_dev, bbuf_complex_dev, permute_staging); if(!gpu_trans) { bbuf_complex = binter_buf; @@ -443,11 +443,11 @@ void block_multiply( gemm_wrapper(hw, thandle, AR, BR, B, M, N, K, alpha, beta, ainter_buf, ainter_buf_dev, bbuf_complex, bbuf_complex_dev, cinter_buf, cinter_tmp_buf_dev); - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, cinter_tmp_buf_dev, is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gemm_wrapper() and transpose_output() enqueue stream-async work on CUDA/HIP; + // gemm_wrapper() and permute_output() enqueue stream-async work on CUDA/HIP; // sync before returning this device block to the pool, which makes the address // immediately re-allocatable. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -466,10 +466,10 @@ void block_multiply( std::span bbuf_real_dev_span = allocate_device_buffer(hw, bsize.value()); T1* bbuf_real_dev = bbuf_real_dev_span.data(); - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf_real, bsize.value(), bdims, blabels, - ainter_buf_dev, bbuf_real_dev, transpose_staging); + ainter_buf_dev, bbuf_real_dev, permute_staging); if(!gpu_trans) { bbuf_real = binter_buf; @@ -479,11 +479,11 @@ void block_multiply( gemm_wrapper(hw, thandle, AR, BR, B, M, N, K, alpha, beta, ainter_buf, ainter_buf_dev, bbuf_real, bbuf_real_dev, cinter_buf, cinter_tmp_buf_dev); - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, cinter_tmp_buf_dev, is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gemm_wrapper() and transpose_output() enqueue stream-async work on CUDA/HIP; + // gemm_wrapper() and permute_output() enqueue stream-async work on CUDA/HIP; // sync before returning this device block to the pool, which makes the address // immediately re-allocatable. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -513,10 +513,10 @@ void block_multiply( std::span abuf_complex_dev_span = allocate_device_buffer(hw, asize.value()); T1* abuf_complex_dev = abuf_complex_dev_span.data(); - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf_complex, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, - blabels, abuf_complex_dev, binter_buf_dev, transpose_staging); + blabels, abuf_complex_dev, binter_buf_dev, permute_staging); if(!gpu_trans) { abuf_complex = ainter_buf; @@ -527,11 +527,11 @@ void block_multiply( gemm_wrapper(hw, thandle, AR, BR, B, M, N, K, alpha, beta, abuf_complex, abuf_complex_dev, binter_buf, binter_buf_dev, cinter_buf, cinter_tmp_buf_dev); - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, cinter_tmp_buf_dev, is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gemm_wrapper() and transpose_output() enqueue stream-async work on CUDA/HIP; + // gemm_wrapper() and permute_output() enqueue stream-async work on CUDA/HIP; // sync before returning this device block to the pool, which makes the address // immediately re-allocatable. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -550,10 +550,10 @@ void block_multiply( std::span abuf_real_dev_span = allocate_device_buffer(hw, asize.value()); T1* abuf_real_dev = abuf_real_dev_span.data(); - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf_real, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf_real, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, - abuf_real_dev, binter_buf_dev, transpose_staging); + abuf_real_dev, binter_buf_dev, permute_staging); if(!gpu_trans) { abuf_real = ainter_buf; @@ -563,11 +563,11 @@ void block_multiply( gemm_wrapper(hw, thandle, AR, BR, B, M, N, K, alpha, beta, abuf_real, abuf_real_dev, binter_buf, binter_buf_dev, cinter_buf, cinter_tmp_buf_dev); - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, cinter_tmp_buf_dev, is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gemm_wrapper() and transpose_output() enqueue stream-async work on CUDA/HIP; + // gemm_wrapper() and permute_output() enqueue stream-async work on CUDA/HIP; // sync before returning this device block to the pool, which makes the address // immediately re-allocatable. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -598,10 +598,10 @@ void block_multiply( gpuMemsetAsync(cbuf_tmp_real_dev, csize.value() * sizeof(T2), thandle); #endif - gpu_trans = transpose_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, + gpu_trans = permute_inputs(hw, thandle, ainter_buf, ainter_dims, ainter_labels, abuf, asize.value(), adims, alabels, binter_buf, binter_dims, binter_labels, bbuf, bsize.value(), bdims, blabels, - ainter_buf_dev, binter_buf_dev, transpose_staging); + ainter_buf_dev, binter_buf_dev, permute_staging); if(!gpu_trans) { copy_data_to_gpu(hw, thandle, ainter_buf, asize.value(), ainter_buf_dev, binter_buf, @@ -619,12 +619,12 @@ void block_multiply( } else { std::copy(cinter_buf_real, cinter_buf_real + csize.value(), cinter_buf); } - transpose_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, + permute_output(hw, thandle, gpu_trans, cinter_buf, cinter_dims, cinter_labels, cbuf, cdims, clabels, cinter_buf_dev, reinterpret_cast(cinter_tmp_buf_dev), is_assign); #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) - // gpu::axpy and the reorder transpose in transpose_output are stream-async on + // gpu::axpy and the permute in permute_output are stream-async on // CUDA/HIP; sync before returning this device block to the pool, which would // otherwise make it immediately re-allocatable while that work is still reading it. if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } @@ -647,7 +647,7 @@ void block_multiply( // cinter_*_dev buffers) and every device buffer allocated internally is free of // in-flight work by the time this function returns. // - // On all backends the gemm/transpose work above is enqueued on `thandle` and + // On all backends the gemm/permute work above is enqueued on `thandle` and // returns immediately (CUDA/HIP: async BLAS; DPC++: in-order queue, no waits // inside gpu::gemm by design), whereas the memory pool's deallocate() is pure // host-side bookkeeping that makes a block instantly re-allocatable -- it does @@ -659,12 +659,12 @@ void block_multiply( // reading it. This device-wide drain per block op is the price of the // immediate-free pool; removing it requires moving the device pool to a // stream-ordered resource (see mr/stream_ordered_memory_resource.hpp). - // The transpose-input staging is returned here, covered by the same drain - // (on the in-order queue the GEMM/output-transpose above are ordered - // behind the input transposes by construction), so each call now pays a + // The permute-input staging is returned here, covered by the same drain + // (on the in-order queue the GEMM/output-permute above are ordered + // behind the input permutes by construction), so each call now pays a // single device-wide sync instead of two. - free_device_buffer(hw, transpose_staging.a); - free_device_buffer(hw, transpose_staging.b); + free_device_buffer(hw, permute_staging.a); + free_device_buffer(hw, permute_staging.b); if(hw == ExecutionHW::GPU) { gpuStreamSynchronize(thandle); } #endif diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index 32d0c5009..ea5d1dd13 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -28,7 +28,7 @@ template void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); -// Out-of-place N-dimensional axis-permuting transpose (in-house reorder +// Out-of-place N-dimensional axis-permuting transpose (in-house permute // kernel; fully replaces the old librett dependency). // out : destination buffer (device), sized as the permuted tensor // in : source buffer (device); must not alias out @@ -40,8 +40,8 @@ void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, // Enqueued on handle.first (cuda/hip stream, in-order SYCL queue); async, // no internal synchronization. template -void transpose_reorder(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, - T scale, bool accumulate, gpuStream_t& handle); +void permute(T* out, const T* in, int ndim, const size_t* outDims, const int* perm, T scale, + bool accumulate, gpuStream_t& handle); } // namespace gpu #endif diff --git a/src/tamm/scheduler.hpp b/src/tamm/scheduler.hpp index bad59bf90..c824b2f5c 100644 --- a/src/tamm/scheduler.hpp +++ b/src/tamm/scheduler.hpp @@ -6,7 +6,7 @@ #include "ga/ga-mpi.h" #include "tamm/dag_impl.hpp" #include "tamm/execution_context.hpp" -#include "tamm/ip_reorder.hpp" +#include "tamm/ip_permute.hpp" #include "tamm/ops.hpp" #include "tamm/tensor.hpp" diff --git a/tests/tamm/Test_TransposeReorder.cpp b/tests/tamm/Test_TransposePermute.cpp similarity index 86% rename from tests/tamm/Test_TransposeReorder.cpp rename to tests/tamm/Test_TransposePermute.cpp index 0a991739e..2ea626d64 100644 --- a/tests/tamm/Test_TransposeReorder.cpp +++ b/tests/tamm/Test_TransposePermute.cpp @@ -1,6 +1,6 @@ -// Unit tests for the in-house GPU tensor transpose (reorder). +// Unit tests for the in-house GPU tensor permute. // -// Exercises the backend-agnostic logic in tamm/kernels/gpu_reorder.hpp +// Exercises the backend-agnostic logic in tamm/kernels/gpu_permute.hpp // (spec building, metadata, index decode, scale/accumulate apply) directly on // the host, so no GPU is needed. Every test compares the exact helpers the // CUDA/HIP/SYCL kernels call against an independent row-major reference for @@ -11,10 +11,10 @@ #define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN #include "doctest/doctest.h" -#include +#include #include -#include -#include +#include +#include #include #include @@ -66,22 +66,22 @@ void reference_transpose(const std::vector& src, std::vector& dst, } } -// Replica of the device loop in gpu_reorder.cpp, but on the host, using the -// shipped helpers (build_reorder_spec / reorder_build_meta / -// reorder_src_index / reorder_scaled / reorder_add). This is the exact math +// Replica of the device loop in gpu_permute.cpp, but on the host, using the +// shipped helpers (build_permute_spec / permute_build_meta / +// permute_src_index / permute_scaled / permute_add). This is the exact math // each backend kernel executes; only the thread loop itself is replicated. template void device_math_replica(const std::vector& src, std::vector& dst, const SizeVec& sdims, const IntLabelVec& slabels, const IntLabelVec& dlabels, T scale, bool accum) { const int ndim = static_cast(sdims.size()); - ReorderSpec spec{}; - build_reorder_spec(sdims, slabels, dlabels, spec); - const ReorderMeta meta = reorder_build_meta(ndim, spec.outDims, spec.perm); - const size_t total = reorder_total(ndim, spec.outDims); + PermuteSpec spec{}; + build_permute_spec(sdims, slabels, dlabels, spec); + const PermuteMeta meta = permute_build_meta(ndim, spec.outDims, spec.perm); + const size_t total = permute_total(ndim, spec.outDims); REQUIRE(dst.size() == total); double sre, sim; - if constexpr(reorder_is_complex_v) { + if constexpr(permute_is_complex_v) { sre = static_cast(scale.real()); sim = static_cast(scale.imag()); } @@ -91,15 +91,15 @@ void device_math_replica(const std::vector& src, std::vector& dst, const S } for(size_t t = 0; t < total; ++t) { const Idx tid = static_cast(t); - const size_t s = reorder_src_index(tid, meta); - const T y = reorder_scaled(src[s], sre, sim); - dst[t] = accum ? reorder_add(dst[t], y) : y; + const size_t s = permute_src_index(tid, meta); + const T y = permute_scaled(src[s], sre, sim); + dst[t] = accum ? permute_add(dst[t], y) : y; } } template T make_val(size_t i) { - if constexpr(reorder_is_complex_v) { + if constexpr(permute_is_complex_v) { return T(static_cast(0.25 * (i + 1)), static_cast(-0.125 * (i + 3))); } @@ -158,7 +158,7 @@ void check_all_perms(const std::vector& sdims, T scale, bool accum) { } // namespace -TEST_CASE("reorder rank 0..2, all perms, scales, assign/accumulate") { +TEST_CASE("permute rank 0..2, all perms, scales, assign/accumulate") { using C = std::complex; for(bool accum: {false, true}) { for(double s: {1.0, 2.5, 0.0, -1.0}) { @@ -173,7 +173,7 @@ TEST_CASE("reorder rank 0..2, all perms, scales, assign/accumulate") { } } -TEST_CASE("reorder rank 3, all perms") { +TEST_CASE("permute rank 3, all perms") { using C = std::complex; for(bool accum: {false, true}) { check_all_perms({2, 3, 4}, 1.0, accum); @@ -184,7 +184,7 @@ TEST_CASE("reorder rank 3, all perms") { } } -TEST_CASE("reorder rank 4, all perms") { +TEST_CASE("permute rank 4, all perms") { using C = std::complex; for(bool accum: {false, true}) { check_all_perms({2, 3, 2, 4}, 1.0, accum); @@ -193,7 +193,7 @@ TEST_CASE("reorder rank 4, all perms") { } } -TEST_CASE("reorder rank 5, all perms") { +TEST_CASE("permute rank 5, all perms") { for(bool accum: {false, true}) { check_all_perms({2, 1, 3, 2, 2}, 1.0, accum); check_all_perms>({2, 1, 2, 2, 3}, std::complex(1.0, 1.0), @@ -201,7 +201,7 @@ TEST_CASE("reorder rank 5, all perms") { } } -TEST_CASE("reorder rank 6 and 8 spot checks") { +TEST_CASE("permute rank 6 and 8 spot checks") { // 720 perms at rank 6: check identity, reversal, and a rotation plus a few // random perms rather than all of them. const std::vector d6{2, 3, 1, 4, 2, 2}; @@ -220,11 +220,11 @@ TEST_CASE("reorder rank 6 and 8 spot checks") { check_case(d8, id8, {3, 7, 1, 5, 0, 6, 2, 4}, 2.0, true); } -TEST_CASE("reorder metadata invariants") { +TEST_CASE("permute metadata invariants") { // Column-major strides + perm round-trip on a known case. const size_t outDims[3] = {4, 2, 3}; const int perm[3] = {2, 0, 1}; - const ReorderMeta meta = reorder_build_meta(3, outDims, perm); + const PermuteMeta meta = permute_build_meta(3, outDims, perm); CHECK(meta.outStrides[0] == 1); CHECK(meta.outStrides[1] == 4); CHECK(meta.outStrides[2] == 8); @@ -232,20 +232,20 @@ TEST_CASE("reorder metadata invariants") { CHECK(meta.inStrides[0] == 1); CHECK(meta.inStrides[1] == 2); CHECK(meta.inStrides[2] == 6); - CHECK(reorder_total(3, outDims) == 24); - CHECK_FALSE(reorder_is_identity(meta)); + CHECK(permute_total(3, outDims) == 24); + CHECK_FALSE(permute_is_identity(meta)); const int idp[3] = {0, 1, 2}; - const ReorderMeta idm = reorder_build_meta(3, outDims, idp); - CHECK(reorder_is_identity(idm)); - CHECK(reorder_meta_fits32(meta, 24)); - CHECK_FALSE(reorder_meta_fits32(meta, size_t{1} << 33)); + const PermuteMeta idm = permute_build_meta(3, outDims, idp); + CHECK(permute_is_identity(idm)); + CHECK(permute_meta_fits32(meta, 24)); + CHECK_FALSE(permute_meta_fits32(meta, size_t{1} << 33)); } // --------------------------------------------------------------------------- -// CPU reorder kernel (HPTT replacement): same natural-order convention as the +// CPU permute kernel (HPTT replacement): same natural-order convention as the // reference above, so no reversal is needed. Covers general alpha/beta // (HPTT supported arbitrary output scaling), both entry points -// (kernels::internal::ip_reorder and blockops::reorder::index_permute_reorder). +// (kernels::internal::ip_permute and blockops::permute::index_permute). // --------------------------------------------------------------------------- using namespace tamm::kernels::cpu; @@ -311,7 +311,7 @@ void check_cpu_case(const std::vector& sdims, const std::vector& sl slabels.begin()); } std::copy(init.begin(), init.end(), got.begin()); - transpose_reorder_cpu(got.data(), src.data(), static_cast(ndim), ddims.data(), + tamm::kernels::cpu::permute(got.data(), src.data(), static_cast(ndim), ddims.data(), perm.data(), alpha, beta); for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); } @@ -324,7 +324,7 @@ void check_cpu_case(const std::vector& sdims, const std::vector& sl SizeVec dv; for(auto d: ddims) dv.emplace_back(d); std::copy(init.begin(), init.end(), got.begin()); - tamm::internal::ip_reorder(got.data(), dv, dl, alpha, src.data(), sv, sl, + tamm::internal::ip_permute(got.data(), dv, dl, alpha, src.data(), sv, sl, beta == T{0}); for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); } @@ -338,7 +338,7 @@ void check_cpu_case(const std::vector& sdims, const std::vector& sl perm.push_back(static_cast(j)); } std::copy(init.begin(), init.end(), got.begin()); - tamm::blockops::reorder::index_permute_reorder(beta, got.data(), alpha, src.data(), perm, + tamm::blockops::permute::index_permute(beta, got.data(), alpha, src.data(), perm, sdims); for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); } @@ -358,7 +358,7 @@ void check_cpu_all_perms(const std::vector& sdims, T alpha, T beta) { std::next_permutation(perm.begin(), perm.end())); } -TEST_CASE("cpu reorder rank 0..2, all perms, general alpha/beta") { +TEST_CASE("cpu permute rank 0..2, all perms, general alpha/beta") { using C = std::complex; for(auto beta_d: {0.0, 1.0, 2.5, -0.5}) { for(auto alpha_d: {1.0, 2.5, 0.0, -1.0}) { @@ -373,7 +373,7 @@ TEST_CASE("cpu reorder rank 0..2, all perms, general alpha/beta") { } } -TEST_CASE("cpu reorder rank 3..5, all perms") { +TEST_CASE("cpu permute rank 3..5, all perms") { using C = std::complex; for(auto beta_d: {0.0, 1.0, 3.0}) { check_cpu_all_perms({2, 3, 4}, 1.0, beta_d); @@ -387,7 +387,7 @@ TEST_CASE("cpu reorder rank 3..5, all perms") { } } -TEST_CASE("cpu reorder rank 6 spot checks + metadata invariants") { +TEST_CASE("cpu permute rank 6 spot checks + metadata invariants") { const std::vector d6{2, 3, 1, 4, 2, 2}; const std::vector id6{0, 1, 2, 3, 4, 5}; check_cpu_case(d6, id6, {5, 4, 3, 2, 1, 0}, 1.0, 0.0); @@ -401,7 +401,7 @@ TEST_CASE("cpu reorder rank 6 spot checks + metadata invariants") { // count from the last axis: {2*3, 3, 1} = {6, 3, 1}. const size_t outDims[3] = {4, 2, 3}; const int perm[3] = {2, 0, 1}; - const ReorderMeta meta = reorder_build_meta_rowmajor(3, outDims, perm); + const PermuteMeta meta = permute_build_meta_rowmajor(3, outDims, perm); CHECK(meta.outStrides[0] == 6); CHECK(meta.outStrides[1] == 3); CHECK(meta.outStrides[2] == 1); @@ -409,5 +409,5 @@ TEST_CASE("cpu reorder rank 6 spot checks + metadata invariants") { CHECK(meta.inStrides[0] == 12); CHECK(meta.inStrides[1] == 4); CHECK(meta.inStrides[2] == 1); - CHECK_FALSE(reorder_is_identity(meta)); + CHECK_FALSE(permute_is_identity(meta)); } diff --git a/tests/tamm/test_tamm.cmake b/tests/tamm/test_tamm.cmake index 7f58e0de4..becb1aef5 100644 --- a/tests/tamm/test_tamm.cmake +++ b/tests/tamm/test_tamm.cmake @@ -27,7 +27,7 @@ add_cxx_unit_test(Test_StrongNum) add_cxx_unit_test(Test_BoundVec) add_cxx_unit_test(Test_BlockSpan) add_cxx_unit_test(Test_MemoryPool) -add_cxx_unit_test(Test_TransposeReorder) +add_cxx_unit_test(Test_TransposePermute) # add_mpi_unit_test(Test_ViewTensor 2 "") # add_mpi_unit_test(Test_QR 2 "") From 43794cf56f9b22a95aa36c0e7e11927683733b21 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Thu, 24 Sep 2026 11:17:20 -0500 Subject: [PATCH 08/14] Rename ip_permute.hpp to index_permute.hpp and Test_TransposePermute to Test_IndexPermute --- src/tamm/CMakeLists.txt | 2 +- src/tamm/block_assign_plan.hpp | 2 +- src/tamm/{ip_permute.hpp => index_permute.hpp} | 0 src/tamm/kernels/gpu_permute.hpp | 2 +- src/tamm/scheduler.hpp | 2 +- tests/tamm/{Test_TransposePermute.cpp => Test_IndexPermute.cpp} | 2 +- tests/tamm/test_tamm.cmake | 2 +- 7 files changed, 6 insertions(+), 6 deletions(-) rename src/tamm/{ip_permute.hpp => index_permute.hpp} (100%) rename tests/tamm/{Test_TransposePermute.cpp => Test_IndexPermute.cpp} (99%) diff --git a/src/tamm/CMakeLists.txt b/src/tamm/CMakeLists.txt index 80a69ffe3..40677d377 100644 --- a/src/tamm/CMakeLists.txt +++ b/src/tamm/CMakeLists.txt @@ -102,7 +102,7 @@ set(TAMM_INCLUDES blockops_cpu.hpp perm.hpp iteration.hpp - ip_permute.hpp + index_permute.hpp blockops_blas.hpp tensor_variant.hpp op_executor.hpp diff --git a/src/tamm/block_assign_plan.hpp b/src/tamm/block_assign_plan.hpp index 5e7849767..debe6b9c9 100644 --- a/src/tamm/block_assign_plan.hpp +++ b/src/tamm/block_assign_plan.hpp @@ -6,7 +6,7 @@ #include "tamm/blockops_blas.hpp" #include "tamm/blockops_cpu.hpp" #include "tamm/errors.hpp" -#include "tamm/ip_permute.hpp" +#include "tamm/index_permute.hpp" #include "tamm/types.hpp" /////////////////////////////////////////////////////////////////////////////// diff --git a/src/tamm/ip_permute.hpp b/src/tamm/index_permute.hpp similarity index 100% rename from src/tamm/ip_permute.hpp rename to src/tamm/index_permute.hpp diff --git a/src/tamm/kernels/gpu_permute.hpp b/src/tamm/kernels/gpu_permute.hpp index 6071b3b5a..7fdf07bbb 100644 --- a/src/tamm/kernels/gpu_permute.hpp +++ b/src/tamm/kernels/gpu_permute.hpp @@ -11,7 +11,7 @@ // // The per-backend kernels in gpu_permute.cpp are thin launch wrappers around // these helpers, so all three backends (CUDA/HIP/SYCL) provably implement the -// same math. The CPU unit test (tests/tamm/Test_TransposePermute.cpp) +// same math. The CPU unit test (tests/tamm/Test_IndexPermute.cpp) // exercises these helpers directly, without needing a GPU. // // Layout note: TAMM callers pass dims/labels in reverse (Fortran) order diff --git a/src/tamm/scheduler.hpp b/src/tamm/scheduler.hpp index c824b2f5c..2620b0da7 100644 --- a/src/tamm/scheduler.hpp +++ b/src/tamm/scheduler.hpp @@ -6,7 +6,7 @@ #include "ga/ga-mpi.h" #include "tamm/dag_impl.hpp" #include "tamm/execution_context.hpp" -#include "tamm/ip_permute.hpp" +#include "tamm/index_permute.hpp" #include "tamm/ops.hpp" #include "tamm/tensor.hpp" diff --git a/tests/tamm/Test_TransposePermute.cpp b/tests/tamm/Test_IndexPermute.cpp similarity index 99% rename from tests/tamm/Test_TransposePermute.cpp rename to tests/tamm/Test_IndexPermute.cpp index 2ea626d64..807eed13c 100644 --- a/tests/tamm/Test_TransposePermute.cpp +++ b/tests/tamm/Test_IndexPermute.cpp @@ -11,7 +11,7 @@ #define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN #include "doctest/doctest.h" -#include +#include #include #include #include diff --git a/tests/tamm/test_tamm.cmake b/tests/tamm/test_tamm.cmake index becb1aef5..b447c366e 100644 --- a/tests/tamm/test_tamm.cmake +++ b/tests/tamm/test_tamm.cmake @@ -27,7 +27,7 @@ add_cxx_unit_test(Test_StrongNum) add_cxx_unit_test(Test_BoundVec) add_cxx_unit_test(Test_BlockSpan) add_cxx_unit_test(Test_MemoryPool) -add_cxx_unit_test(Test_TransposePermute) +add_cxx_unit_test(Test_IndexPermute) # add_mpi_unit_test(Test_ViewTensor 2 "") # add_mpi_unit_test(Test_QR 2 "") From 006a43b4da7af41c313d73c079590d3c68625003 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Thu, 24 Sep 2026 11:20:09 -0500 Subject: [PATCH 09/14] Rename internal ip_permute to index_permute overload --- src/tamm/kernels/assign.hpp | 6 +++--- tests/tamm/Test_IndexPermute.cpp | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/tamm/kernels/assign.hpp b/src/tamm/kernels/assign.hpp index 1c870c429..55b3573dd 100644 --- a/src/tamm/kernels/assign.hpp +++ b/src/tamm/kernels/assign.hpp @@ -373,8 +373,8 @@ void ip_gen_loop(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T sca } template -void ip_permute(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, - const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign = true) { +void index_permute(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, const T* src, + const SizeVec& sdims, const IntLabelVec& slabels, bool is_assign = true) { const size_t ndim = ddims.size(); EXPECTS(sdims.size() == ndim && dlabels.size() == ndim && slabels.size() == ndim); EXPECTS(ndim <= static_cast(kernels::gpu::permute_maxrank)); @@ -473,7 +473,7 @@ void assign(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T scale, c if(is_assign) { internal::index_permute(dst, src, perm_to_dest, ddims, scale); } else { internal::index_permute_acc(dst, src, perm_to_dest, ddims, scale); } } - else internal::ip_permute(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); + else internal::index_permute(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } else { internal::ip_gen_loop(dst, ddims, dlabels, scale, src, sdims, slabels, is_assign); } } diff --git a/tests/tamm/Test_IndexPermute.cpp b/tests/tamm/Test_IndexPermute.cpp index 807eed13c..90c7c4d21 100644 --- a/tests/tamm/Test_IndexPermute.cpp +++ b/tests/tamm/Test_IndexPermute.cpp @@ -245,7 +245,7 @@ TEST_CASE("permute metadata invariants") { // CPU permute kernel (HPTT replacement): same natural-order convention as the // reference above, so no reversal is needed. Covers general alpha/beta // (HPTT supported arbitrary output scaling), both entry points -// (kernels::internal::ip_permute and blockops::permute::index_permute). +// (internal::index_permute and blockops::permute::index_permute). // --------------------------------------------------------------------------- using namespace tamm::kernels::cpu; @@ -324,7 +324,7 @@ void check_cpu_case(const std::vector& sdims, const std::vector& sl SizeVec dv; for(auto d: ddims) dv.emplace_back(d); std::copy(init.begin(), init.end(), got.begin()); - tamm::internal::ip_permute(got.data(), dv, dl, alpha, src.data(), sv, sl, + tamm::internal::index_permute(got.data(), dv, dl, alpha, src.data(), sv, sl, beta == T{0}); for(size_t i = 0; i < total; ++i) CHECK(got[i] == ref[i]); } From 84f5a2d4f7cd198806deba0504781bf05d543781 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty <59661409+abagusetty@users.noreply.github.com> Date: Mon, 28 Sep 2026 12:17:50 -0500 Subject: [PATCH 10/14] Update hip memory allocation function --- src/tamm/mr/pinned_memory_resource.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tamm/mr/pinned_memory_resource.hpp b/src/tamm/mr/pinned_memory_resource.hpp index d1b888e73..d91be0dd5 100644 --- a/src/tamm/mr/pinned_memory_resource.hpp +++ b/src/tamm/mr/pinned_memory_resource.hpp @@ -67,7 +67,7 @@ class pinned_memory_resource final: public host_memory_resource { auto status = cudaMallocHost(&ptr, size); if(cudaSuccess != status) { throw std::bad_alloc{}; } #elif defined(USE_HIP) - auto status = hipMallocHost(&ptr, size); + auto status = hipHostMalloc(&ptr, size); if (hipSuccess != status) { throw std::bad_alloc{}; } #elif defined(USE_DPCPP) // aligned_alloc_host, not malloc_host -- see the note in gpu_memory_resource.hpp. From edffa99457997a389e36f3d9d44c793050534d18 Mon Sep 17 00:00:00 2001 From: Ajay Panyala Date: Wed, 30 Sep 2026 08:34:28 -0700 Subject: [PATCH 11/14] fix merge error in tamm_blas.hpp --- docs/install.rst | 2 +- src/tamm/kernels/gpu_blas.cpp | 6 +++--- src/tamm/kernels/tamm_blas.hpp | 6 +++++- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/install.rst b/docs/install.rst index c84d74f50..44d26b90c 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -202,7 +202,7 @@ Polaris modules and env module use /soft/modulefiles/ module load PrgEnv-gnu - module load cudatoolkit-standalone/12.6.1 spack-pe-base cmake + module load cudatoolkit-standalone spack-pe-base cmake module unload craype-accel-nvidia80 export CRAYPE_LINK_TYPE=dynamic diff --git a/src/tamm/kernels/gpu_blas.cpp b/src/tamm/kernels/gpu_blas.cpp index 9d7794499..23e3fb399 100644 --- a/src/tamm/kernels/gpu_blas.cpp +++ b/src/tamm/kernels/gpu_blas.cpp @@ -44,7 +44,7 @@ void tamm::kernels::gpu::axpy(const int64_t n, const T* src, const int incx, T*& template void tamm::kernels::gpu::gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, - int lda, const T beta, T1* C, int ldc, gpuStream_t& handle) { + int lda, const T beta, T1* C, int ldc, gpuStream_t& handle) { #if defined(USE_DPCPP) #ifdef USE_PORT_BLAS blas::SB_Handle sb_handle(handle.first); @@ -52,8 +52,8 @@ void tamm::kernels::gpu::gemm(int n, int m, int k, const T alpha, const T3* B, i const_cast(A), lda, beta, C, ldc, {}); #else oneapi::mkl::blas::column_major::gemm(handle.first, oneapi::mkl::transpose::N, - oneapi::mkl::transpose::N, n, m, k, alpha, B, ldb, A, - lda, beta, C, ldc); + oneapi::mkl::transpose::N, n, m, k, alpha, B, ldb, A, lda, + beta, C, ldc); #endif // USE_PORT_BLAS #elif defined(USE_CUDA) diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index e4adfd70c..79df2f619 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -25,6 +25,10 @@ template void axpy(const int64_t n, const T* src, const int incx, T*& dst, const int incy, gpuStream_t& gpuhandle); +template +void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, int lda, + const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); + // Out-of-place N-dimensional axis-permuting transpose (in-house permute // kernel; fully replaces the old librett dependency). // out : destination buffer (device), sized as the permuted tensor @@ -53,7 +57,7 @@ void gemm(blas::Op transa, blas::Op transb, int m, int n, int k, const T alpha, template void gesvd(lapack::Job jobu, lapack::Job jobvt, int64_t m, int64_t n, T* A, int64_t lda, blas::real_type* S, T* U, int64_t ldu, T* VT, int64_t ldvt); - + } // namespace gpu #endif From a86ee59605965771502b7ad105541d367ff1a35a Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty <59661409+abagusetty@users.noreply.github.com> Date: Wed, 30 Sep 2026 10:46:26 -0500 Subject: [PATCH 12/14] Update gemm function documentation in tamm_blas.hpp Added a note regarding the generalized eigensolver API and a TODO for uniform gemm() API usage. --- src/tamm/kernels/tamm_blas.hpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index 79df2f619..f2718b7f2 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -46,6 +46,8 @@ void permute(T* out, const T* in, int ndim, const size_t* outDims, const int* pe // Column-major C = alpha op(A) op(B) + beta C on device pointers, with the same argument order // and meaning as blas::gemm(blas::Layout::ColMajor, ...). Runs on gpuhandle's stream. +// Note: Added for the generalized eigensolver API that needs to call gemms on 2D mats with transpose. +// TODO: Address uniform gemm() API usage template void gemm(blas::Op transa, blas::Op transb, int m, int n, int k, const T alpha, const T* A, int lda, const T* B, int ldb, const T beta, T* C, int ldc, gpuStream_t& gpuhandle); From 3603bc366c122ddc55887e7f4dc9199cfd5794fe Mon Sep 17 00:00:00 2001 From: Ajay Panyala Date: Wed, 30 Sep 2026 09:46:11 -0700 Subject: [PATCH 13/14] update install.rst --- docs/install.rst | 7 +++---- docs/prerequisites.rst | 32 -------------------------------- 2 files changed, 3 insertions(+), 36 deletions(-) delete mode 100644 docs/prerequisites.rst diff --git a/docs/install.rst b/docs/install.rst index 44d26b90c..eae4630df 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -1,5 +1,3 @@ -The prerequisites needed to build this repository can be found :doc:`here `. - Build Instructions ================== @@ -10,7 +8,7 @@ Dependencies * cmake >= 3.26 * MPI -* C++17 compiler (information on supported compilers here :doc:`here `.) +* C++20 Compilers (GNU >= 14.1, LLVM Clang >= 18, AppleClang >= 19) * CUDA >= 12.8 (Required only for CUDA builds) * ROCM >= 6.2 (Required only for ROCM builds) @@ -123,7 +121,8 @@ Default build on MACOS ~~~~~~~~~~~~~~~~~~~~~~ .. note:: - The prerequisites for ``MACOS`` can be installed using ``brew`` as detailed :doc:`here `. + The required external dependencies for ``MACOS`` can be installed using ``brew``: + ``brew install gcc mpich cmake autoconf automake libtool`` :: diff --git a/docs/prerequisites.rst b/docs/prerequisites.rst deleted file mode 100644 index eaa6fc036..000000000 --- a/docs/prerequisites.rst +++ /dev/null @@ -1,32 +0,0 @@ -:orphan: - -Supported Compilers -=================== - -- GCC versions >= 9.1 -- LLVM Clang >= 9 - -Supported Configurations -======================== - -- The following configurations are recommended since they are tested - and are known to work: - - - GCC versions >= 9.1 + OpenMPI-2.x/MPICH-3.x built using the - corresponding GCC versions. - - LLVM Clang versions >= 9 + OpenMPI-2.x/MPICH-3.x - -Installing Prerequisites -======================== - -On Mac OSX -=========== - -- brew install gcc mpich cmake autoconf automake libtool - -On Linux -========= - -We recommend using the `Spack package manager `__ to -install and manage the prerequisites if the versions avaiable via the OS -package manager are not sufficient. From f389fa0dd5e14e636c5dda547535634ccf829650 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Wed, 30 Sep 2026 12:35:39 -0500 Subject: [PATCH 14/14] cleanup librett and hptt strings in the rest of the code --- src/tamm/index_permute.hpp | 5 ++--- src/tamm/kernels/assign.hpp | 4 ++-- src/tamm/kernels/cpu_permute.hpp | 2 +- src/tamm/kernels/gpu_permute.hpp | 7 +++++-- src/tamm/kernels/multiply.hpp | 3 +-- src/tamm/kernels/tamm_blas.hpp | 2 +- src/tamm/scheduler.hpp | 11 ++++++----- tests/tamm/Test_IndexPermute.cpp | 4 ++-- 8 files changed, 20 insertions(+), 18 deletions(-) diff --git a/src/tamm/index_permute.hpp b/src/tamm/index_permute.hpp index 926da53e0..bcc6527ef 100644 --- a/src/tamm/index_permute.hpp +++ b/src/tamm/index_permute.hpp @@ -12,17 +12,16 @@ namespace tamm::blockops::permute { /////////////////////////////////////////////////////////////////////////////// // // Index permute using the in-house CPU permute kernel -// (replaces the old HPTT-based index_permute_hptt) // /////////////////////////////////////////////////////////////////////////////// -// Out-of-place permute with HPTT-identical semantics: +// Out-of-place permute with assign/accumulate semantics: // lbuf = lscale * lbuf + rscale * permute(rbuf) // perm_to_dest maps each destination axis to its source axis and sdims holds // the source extents, both in natural order; the output extents follow as // outDims[k] == sdims[perm_to_dest[k]]. Scale types are independent template // parameters so call sites may mix e.g. an int literal rscale with a -// complex buffer (as the old single-T HPTT wrapper allowed via conversion). +// complex buffer. template void index_permute(BL lscale, T1* lbuf, BA rscale, const T2* rbuf, const PermVector& perm_to_dest, const std::vector& sdims) { diff --git a/src/tamm/kernels/assign.hpp b/src/tamm/kernels/assign.hpp index 55b3573dd..7ba92d82c 100644 --- a/src/tamm/kernels/assign.hpp +++ b/src/tamm/kernels/assign.hpp @@ -379,8 +379,8 @@ void index_permute(T* dst, const SizeVec& ddims, const IntLabelVec& dlabels, T s EXPECTS(sdims.size() == ndim && dlabels.size() == ndim && slabels.size() == ndim); EXPECTS(ndim <= static_cast(kernels::gpu::permute_maxrank)); - // Natural-order output extents + output-axis -> source-axis map, exactly - // the (size, perm) pair the old HPTT plan consumed. + // Natural-order output extents + output-axis -> source-axis map, the same + // (size, perm) convention the CPU and GPU permute kernels consume. size_t outDims[kernels::gpu::permute_maxrank] = {}; int perm[kernels::gpu::permute_maxrank] = {}; for(size_t i = 0; i < ndim; i++) { diff --git a/src/tamm/kernels/cpu_permute.hpp b/src/tamm/kernels/cpu_permute.hpp index eada57b0e..83afd80e4 100644 --- a/src/tamm/kernels/cpu_permute.hpp +++ b/src/tamm/kernels/cpu_permute.hpp @@ -1,6 +1,6 @@ #pragma once -// Host-only CPU tensor permute, the HPTT replacement. +// Host-only CPU tensor permute. // // Row-major (TAMM-native) counterpart of the GPU kernel in gpu_permute.hpp. // Reuses its PermuteMeta storage, elementwise scale/add helpers and identity diff --git a/src/tamm/kernels/gpu_permute.hpp b/src/tamm/kernels/gpu_permute.hpp index 7fdf07bbb..2b4c1b8d2 100644 --- a/src/tamm/kernels/gpu_permute.hpp +++ b/src/tamm/kernels/gpu_permute.hpp @@ -5,7 +5,8 @@ // This header holds everything about the permute that is backend-agnostic: // - the permutation metadata layout, // - the output-linear-index -> input-linear-index decode (column-major, -// first-axis-fastest, matching the convention the old librett path used), +// first-axis-fastest; callers pass dims/labels in reversed Fortran order, +// so this implements the natural row-major transpose), // - the scale/accumulate elementwise apply, // - host-side helpers that build the metadata from TAMM's (dims, labels). // @@ -200,7 +201,9 @@ inline void build_permute_spec(const SizeVec& sdims, const IntLabelVec& slabels, EXPECTS(dlabels.size() == ndim && slabels.size() == ndim); EXPECTS(ndim <= static_cast(permute_maxrank)); - // Reverse into Fortran order (librett used the same convention). + // Reverse into Fortran order for the column-major kernel: a row-major + // transpose with perm P equals a column-major transpose with the reversed + // perm, so reversing both dim orders converts between the two layouts. size_t r_size[permute_maxrank]; IntLabel r_slabels[permute_maxrank]; IntLabel r_dlabels[permute_maxrank]; diff --git a/src/tamm/kernels/multiply.hpp b/src/tamm/kernels/multiply.hpp index eb2e29eaa..cdae91be6 100644 --- a/src/tamm/kernels/multiply.hpp +++ b/src/tamm/kernels/multiply.hpp @@ -216,8 +216,7 @@ void permute_output(ExecutionHW hw, gpuStream_t& thandle, bool gpu_trans, T1* ci T1*& cinter_tmp_buf_dev, bool is_assign) { #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP) if(hw == ExecutionHW::GPU) { - // NOTE: always overwrite (is_assign=true) on the GPU path, matching the - // old librett assign_gpu semantics (which ignored is_assign/scale). + // NOTE: always overwrite (is_assign=true) on the GPU path. // Accumulation across reduction iterations here is performed by GEMM // itself (beta=cscale=1 accumulates in cinter_tmp_buf_dev in inter // ordering), so the output permute must overwrite cinter_buf_dev with diff --git a/src/tamm/kernels/tamm_blas.hpp b/src/tamm/kernels/tamm_blas.hpp index f2718b7f2..7d9871d5f 100644 --- a/src/tamm/kernels/tamm_blas.hpp +++ b/src/tamm/kernels/tamm_blas.hpp @@ -30,7 +30,7 @@ void gemm(int n, int m, int k, const T alpha, const T3* B, int ldb, const T2* A, const T beta, T1* C, int ldc, gpuStream_t& gpuhandle); // Out-of-place N-dimensional axis-permuting transpose (in-house permute -// kernel; fully replaces the old librett dependency). +// kernel). // out : destination buffer (device), sized as the permuted tensor // in : source buffer (device); must not alias out // ndim : tensor rank (0 .. 8) diff --git a/src/tamm/scheduler.hpp b/src/tamm/scheduler.hpp index 2620b0da7..b8dc9b0b7 100644 --- a/src/tamm/scheduler.hpp +++ b/src/tamm/scheduler.hpp @@ -519,16 +519,16 @@ class Scheduler { // PermVector perm_to_dest{perm.begin(), perm.end()}; // TensorType lscale = is_assign ? TensorType{0} : TensorType{1}; - // auto lambda_hptt = [&](const IndexVector& itval) { + // auto lambda_permute = [&](const IndexVector& itval) { // IndexVector src_id = itval; // size_t dst_size = dst.block_size(itval); // std::vector src_buf(dst_size); // src.get(src_id, src_buf); // std::vector dest_buf(dst_size); - // blockops::hptt::index_permute_hptt(lscale, dest_buf.data(), scale, - // src_buf.data(), perm_to_dest, - // src.block_dims(src_id)); + // blockops::permute::index_permute(lscale, dest_buf.data(), scale, + // src_buf.data(), perm_to_dest, + // src.block_dims(src_id)); // if(is_assign) // dst.put(itval, dest_buf); @@ -537,7 +537,8 @@ class Scheduler { // }; // auto ec = dst.execution_context(); - // block_for(*ec, dst(), lambda_hptt); + // block_for(*ec, dst(), lambda_permute); + // // NOTE: enabling this requires #include "tamm/work.hpp" for block_for. // } private: diff --git a/tests/tamm/Test_IndexPermute.cpp b/tests/tamm/Test_IndexPermute.cpp index 90c7c4d21..6975b6ee9 100644 --- a/tests/tamm/Test_IndexPermute.cpp +++ b/tests/tamm/Test_IndexPermute.cpp @@ -242,9 +242,9 @@ TEST_CASE("permute metadata invariants") { } // --------------------------------------------------------------------------- -// CPU permute kernel (HPTT replacement): same natural-order convention as the +// CPU permute kernel: same natural-order convention as the // reference above, so no reversal is needed. Covers general alpha/beta -// (HPTT supported arbitrary output scaling), both entry points +// output scaling, both entry points // (internal::index_permute and blockops::permute::index_permute). // ---------------------------------------------------------------------------