From facac5cc0d160c395291e6a80f8ec3c05a1e2daf Mon Sep 17 00:00:00 2001 From: Zheming Jin <57232910+zjin-lcf@users.noreply.github.com> Date: Sat, 15 Aug 2026 10:19:35 -0500 Subject: [PATCH 1/3] [blas] Implement omatcopy2 for the cuBLAS and rocBLAS backends omatcopy2 applies an element stride to each matrix, which geam cannot express, so both backends reported it as unimplemented. Add a portable SYCL kernel shared by the two backends: a transposing variant that stages a tile through local memory, so that neither the load nor the store steps by a leading dimension, and a plain strided copy for the nontrans case. Unit strides make omatcopy2 equivalent to omatcopy, but routing that case to geam turns out to cost more than it saves. The vendor libraries have to be driven from a host task and with a stream synchronize, whereas these kernels are ordinary asynchronous SYCL. On gfx950 the kernels beat the geam path by 37-77% even when the caller waits after every call, so the dispatch always uses them. The tile shape is chosen per element size, and separately for unit and non-unit strides: with unit strides both accesses are contiguous and the two phases want equal width, while with real strides widening the stores buys nothing and a taller tile amortises the per-tile overhead. Every variant keeps the tile under 17 KB so that three groups stay resident even on CDNA1-CDNA3, which have a quarter of CDNA4's local memory. Extents were measured on gfx950 over strides 1-4 and sizes from 64 to 8192 square. Tested on gfx950 (MI350X): the omatcopy2 unit tests pass for both the compile-time and run-time dispatch layers, alongside a standalone check of 336 configurations covering both layouts, all three transpose modes, unit and mixed and general strides, and sizes straddling the tile extents. --- .../backends/cublas/cublas_extensions.cpp | 105 +++----- src/blas/backends/omatcopy2_kernels.hpp | 246 ++++++++++++++++++ .../backends/rocblas/rocblas_extensions.cpp | 117 ++++----- 3 files changed, 338 insertions(+), 130 deletions(-) create mode 100644 src/blas/backends/omatcopy2_kernels.hpp diff --git a/src/blas/backends/cublas/cublas_extensions.cpp b/src/blas/backends/cublas/cublas_extensions.cpp index a556d1140..fea864769 100644 --- a/src/blas/backends/cublas/cublas_extensions.cpp +++ b/src/blas/backends/cublas/cublas_extensions.cpp @@ -18,6 +18,7 @@ **************************************************************************/ #include "cublas_helper.hpp" #include "cublas_task.hpp" +#include "blas/backends/omatcopy2_kernels.hpp" #include "oneapi/math/exceptions.hpp" #include "oneapi/math/blas/detail/cublas/onemath_blas_cublas.hpp" @@ -121,25 +122,23 @@ OMATCOPY_LAUNCHER(std::complex, cublasZgeam) #undef OMATCOPY_LAUNCHER -template -void omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, int64_t m, - int64_t n, T alpha, sycl::buffer& a, int64_t lda, std::int64_t stridea, - sycl::buffer& b, int64_t ldb, std::int64_t strideb) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER(TYPE, CUBLAS_ROUTINE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2(#CUBLAS_ROUTINE, CUBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, lda, b, \ - ldb, strideb); \ +// Unit element strides make omatcopy2 equivalent to omatcopy, but routing them +// to geam is not worth it: cuBLAS has to be called from a host task and with a +// stream synchronize, whereas these kernels are plain asynchronous SYCL. The +// equivalent rocBLAS path measured 37-77% slower than the kernels on gfx950 +// even when the caller waits after every call. +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::col_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb); \ } -OMATCOPY2_LAUNCHER(float, "unimplemented") -OMATCOPY2_LAUNCHER(double, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER(float) +OMATCOPY2_LAUNCHER(double) +OMATCOPY2_LAUNCHER(std::complex) +OMATCOPY2_LAUNCHER(std::complex) #undef OMATCOPY2_LAUNCHER void imatcopy(sycl::queue& queue, transpose trans, int64_t m, int64_t n, float alpha, @@ -302,25 +301,19 @@ OMATCOPY_LAUNCHER_USM(std::complex, cublasZgeam) #undef OMATCOPY_LAUNCHER_USM -template -sycl::event omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, - int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, - int64_t ldb, int64_t strideb, const std::vector& dependencies) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER_USM(TYPE, CUBLAS_ROUTINE) \ +#define OMATCOPY2_LAUNCHER_USM(TYPE) \ sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2(#CUBLAS_ROUTINE, CUBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, \ - lda, b, ldb, strideb, dependencies); \ + return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::col_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb, \ + dependencies); \ } -OMATCOPY2_LAUNCHER_USM(float, "unimplemented") -OMATCOPY2_LAUNCHER_USM(double, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER_USM(float) +OMATCOPY2_LAUNCHER_USM(double) +OMATCOPY2_LAUNCHER_USM(std::complex) +OMATCOPY2_LAUNCHER_USM(std::complex) #undef OMATCOPY2_LAUNCHER_USM sycl::event imatcopy(sycl::queue& queue, transpose trans, int64_t m, int64_t n, float alpha, @@ -484,25 +477,19 @@ OMATCOPY_LAUNCHER(std::complex, cublasZgeam) #undef OMATCOPY_LAUNCHER -template -void omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, int64_t m, - int64_t n, T alpha, sycl::buffer& a, int64_t lda, std::int64_t stridea, - sycl::buffer& b, int64_t ldb, std::int64_t strideb) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER(TYPE, CUBLAS_ROUTINE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2(#CUBLAS_ROUTINE, CUBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, lda, b, \ - ldb, strideb); \ +// See the column-major overloads for why geam is not used at unit stride. +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::row_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb); \ } -OMATCOPY2_LAUNCHER(float, "unimplemented") -OMATCOPY2_LAUNCHER(double, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER(float) +OMATCOPY2_LAUNCHER(double) +OMATCOPY2_LAUNCHER(std::complex) +OMATCOPY2_LAUNCHER(std::complex) #undef OMATCOPY2_LAUNCHER void imatcopy(sycl::queue& queue, transpose trans, int64_t m, int64_t n, float alpha, @@ -665,25 +652,19 @@ OMATCOPY_LAUNCHER_USM(std::complex, cublasZgeam) #undef OMATCOPY_LAUNCHER_USM -template -sycl::event omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, - int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, - int64_t ldb, int64_t strideb, const std::vector& dependencies) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER_USM(TYPE, CUBLAS_ROUTINE) \ +#define OMATCOPY2_LAUNCHER_USM(TYPE) \ sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2(#CUBLAS_ROUTINE, CUBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, \ - lda, b, ldb, strideb, dependencies); \ + return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::row_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb, \ + dependencies); \ } -OMATCOPY2_LAUNCHER_USM(float, "unimplemented") -OMATCOPY2_LAUNCHER_USM(double, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER_USM(float) +OMATCOPY2_LAUNCHER_USM(double) +OMATCOPY2_LAUNCHER_USM(std::complex) +OMATCOPY2_LAUNCHER_USM(std::complex) #undef OMATCOPY2_LAUNCHER_USM sycl::event imatcopy(sycl::queue& queue, transpose trans, int64_t m, int64_t n, float alpha, diff --git a/src/blas/backends/omatcopy2_kernels.hpp b/src/blas/backends/omatcopy2_kernels.hpp new file mode 100644 index 000000000..b02a68bc4 --- /dev/null +++ b/src/blas/backends/omatcopy2_kernels.hpp @@ -0,0 +1,246 @@ +/*************************************************************************** +* Copyright (C) Codeplay Software Limited +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* For your convenience, a copy of the License has been included in this +* repository. +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +* +**************************************************************************/ + +/** + * @file omatcopy2_kernels.hpp : portable SYCL kernels implementing the + * omatcopy2 element strides, which vendor geam entry points cannot express. + */ +#ifndef _ONEMATH_BLAS_OMATCOPY2_KERNELS_HPP_ +#define _ONEMATH_BLAS_OMATCOPY2_KERNELS_HPP_ + +#if __has_include() +#include +#else +#include +#endif + +#include +#include +#include +#include + +#include "oneapi/math/types.hpp" + +namespace oneapi { +namespace math { +namespace blas { +namespace omatcopy2_kernels { + +// Geometry of the local-memory tile that stages the transpose. Without the +// tile one of the two global accesses would step by the leading dimension; +// with it both step by the (small) element stride instead. +// +// The tile is `rows` x `cols` elements and the work-group is `block` x `rows` +// items: the load phase puts the lane index on the row extent and the store +// phase puts it on the column extent, so `cols` sets how wide the stores are. +// +// Both variants below keep the tile under 17 KB for every element type, which +// leaves room for three resident groups even on CDNA1-CDNA3 (64 KB of local +// memory per compute unit); CDNA4 has 160 KB. Extents were measured on gfx950 +// over strides 1-4 and sizes from 64 to 8192 square. + +/// Unit element strides: both global accesses are contiguous, so the two +/// phases want the same width and a square tile is best. +template +struct unit_stride_geometry { + static constexpr int rows = sizeof(T) <= 4 ? 64 : 32; + static constexpr int cols = rows; + static constexpr int block = 8; +}; + +/// Real element strides: neither access is contiguous, so widening the stores +/// buys nothing and a taller tile amortises the per-tile overhead instead. +template +struct strided_geometry { + static constexpr int rows = sizeof(T) <= 8 ? 64 : 32; + static constexpr int cols = 32; + static constexpr int block = sizeof(T) == 8 ? 16 : 8; +}; + +template +struct is_complex : std::false_type {}; +template +struct is_complex> : std::true_type {}; + +template +inline T conj_if(const T& value, bool do_conj) { + if constexpr (is_complex::value) { + return do_conj ? T(value.real(), -value.imag()) : value; + } + else { + (void)do_conj; + return value; + } +} + +/// b[r * strideb + c * ldb] = alpha * a[r * stridea + c * lda] +/// +/// The fastest-moving index is r, so both accesses run along the element +/// stride and stay as coalesced as the strides allow. +template +void launch_nontrans(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T alpha, AccessorA a, + int64_t lda, int64_t stridea, AccessorB b, int64_t ldb, int64_t strideb) { + cgh.parallel_for(sycl::range<2>(static_cast(logical_n), static_cast(logical_m)), + [=](sycl::id<2> id) { + const int64_t c = static_cast(id[0]); + const int64_t r = static_cast(id[1]); + b[r * strideb + c * ldb] = alpha * a[r * stridea + c * lda]; + }); +} + +/// b[c * strideb + r * ldb] = alpha * op(a[r * stridea + c * lda]) +/// +/// Loads run along r and stores run along c, both with the lane index as the +/// fastest dimension, so neither global access steps by a leading dimension. +template +void launch_trans(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T alpha, bool do_conj, + AccessorA a, int64_t lda, int64_t stridea, AccessorB b, int64_t ldb, + int64_t strideb) { + // Odd padding keeps the transposed local reads off a single bank. + constexpr int pitch = Geom::rows + 1; + + static_assert(Geom::cols <= Geom::rows, "the group lane count comes from the row extent"); + static_assert(Geom::rows % Geom::block == 0 && Geom::cols % Geom::block == 0, + "both tile extents must be covered by whole steps of block"); + + sycl::local_accessor tile(sycl::range<1>(Geom::cols * pitch), cgh); + + const int64_t tiles_r = (logical_m + Geom::rows - 1) / Geom::rows; + const int64_t tiles_c = (logical_n + Geom::cols - 1) / Geom::cols; + + const sycl::range<2> global(static_cast(tiles_c) * Geom::block, + static_cast(tiles_r) * Geom::rows); + const sycl::range<2> local(Geom::block, Geom::rows); + + cgh.parallel_for(sycl::nd_range<2>(global, local), [=](sycl::nd_item<2> item) { + const int ly = static_cast(item.get_local_id(0)); + const int lx = static_cast(item.get_local_id(1)); + const int64_t tile_r = static_cast(item.get_group(1)) * Geom::rows; + const int64_t tile_c = static_cast(item.get_group(0)) * Geom::cols; + + const int64_t load_r = tile_r + lx; + if (load_r < logical_m) { + for (int k = 0; k < Geom::cols; k += Geom::block) { + const int cl = ly + k; + const int64_t c = tile_c + cl; + if (c < logical_n) { + tile[cl * pitch + lx] = a[load_r * stridea + c * lda]; + } + } + } + + item.barrier(sycl::access::fence_space::local_space); + + // Stores use only the first `cols` lanes; the rest of the group idles. + const int64_t store_c = tile_c + lx; + if (lx < Geom::cols && store_c < logical_n) { + for (int k = 0; k < Geom::rows; k += Geom::block) { + const int rl = ly + k; + const int64_t r = tile_r + rl; + if (r < logical_m) { + b[store_c * strideb + r * ldb] = + alpha * conj_if(tile[lx * pitch + rl], do_conj); + } + } + } + }); +} + +/// Picks the tile shape from the strides, which decide whether the global +/// accesses are contiguous. +template +void launch_trans_dispatch(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T alpha, + bool do_conj, AccessorA a, int64_t lda, int64_t stridea, AccessorB b, + int64_t ldb, int64_t strideb) { + if (stridea == 1 && strideb == 1) { + launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, lda, + stridea, b, ldb, strideb); + } + else { + launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, lda, + stridea, b, ldb, strideb); + } +} + +template +sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, + oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, const T* a, + int64_t lda, int64_t stridea, T* b, int64_t ldb, int64_t strideb, + const std::vector& dependencies) { + const bool col_major = (layout == oneapi::math::layout::col_major); + const int64_t logical_m = col_major ? m : n; + const int64_t logical_n = col_major ? n : m; + const bool do_trans = (trans != oneapi::math::transpose::nontrans); + const bool do_conj = (trans == oneapi::math::transpose::conjtrans); + + if (logical_m <= 0 || logical_n <= 0) { + return queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + cgh.single_task([=]() {}); + }); + } + + return queue.submit([&](sycl::handler& cgh) { + cgh.depends_on(dependencies); + if (do_trans) { + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a, lda, stridea, b, + ldb, strideb); + } + else { + launch_nontrans(cgh, logical_m, logical_n, alpha, a, lda, stridea, b, ldb, strideb); + } + }); +} + +template +void omatcopy2_buffer(sycl::queue& queue, oneapi::math::layout layout, + oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, + sycl::buffer& a, int64_t lda, int64_t stridea, sycl::buffer& b, + int64_t ldb, int64_t strideb) { + const bool col_major = (layout == oneapi::math::layout::col_major); + const int64_t logical_m = col_major ? m : n; + const int64_t logical_n = col_major ? n : m; + const bool do_trans = (trans != oneapi::math::transpose::nontrans); + const bool do_conj = (trans == oneapi::math::transpose::conjtrans); + + if (logical_m <= 0 || logical_n <= 0) { + return; + } + + queue.submit([&](sycl::handler& cgh) { + auto a_acc = a.template get_access(cgh); + // Strided writes skip elements, so the untouched ones must be preserved. + auto b_acc = b.template get_access(cgh); + if (do_trans) { + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a_acc, lda, stridea, + b_acc, ldb, strideb); + } + else { + launch_nontrans(cgh, logical_m, logical_n, alpha, a_acc, lda, stridea, b_acc, ldb, + strideb); + } + }); +} + +} // namespace omatcopy2_kernels +} // namespace blas +} // namespace math +} // namespace oneapi + +#endif // _ONEMATH_BLAS_OMATCOPY2_KERNELS_HPP_ diff --git a/src/blas/backends/rocblas/rocblas_extensions.cpp b/src/blas/backends/rocblas/rocblas_extensions.cpp index b649ed1ef..a347c103b 100644 --- a/src/blas/backends/rocblas/rocblas_extensions.cpp +++ b/src/blas/backends/rocblas/rocblas_extensions.cpp @@ -22,6 +22,7 @@ #include "rocblas_helper.hpp" #include "rocblas_task.hpp" +#include "blas/backends/omatcopy2_kernels.hpp" #include "oneapi/math/exceptions.hpp" #include "oneapi/math/blas/detail/rocblas/onemath_blas_rocblas.hpp" @@ -130,25 +131,23 @@ OMATCOPY_LAUNCHER(std::complex, rocblas_zgeam) #undef OMATCOPY_LAUNCHER -template -void omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, int64_t m, - int64_t n, T alpha, sycl::buffer& a, int64_t lda, std::int64_t stridea, - sycl::buffer& b, int64_t ldb, std::int64_t strideb) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER(TYPE, ROCBLAS_ROUTINE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2(#ROCBLAS_ROUTINE, ROCBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, lda, \ - b, ldb, strideb); \ +// Unit element strides make omatcopy2 equivalent to omatcopy, but routing them +// to geam is not worth it: rocBLAS has to be called from a host task and with a +// stream synchronize, whereas these kernels are plain asynchronous SYCL. On +// gfx950 the kernels beat the geam path by 37-77% even when the caller waits +// after every call, and by more when calls can overlap. +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::col_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb); \ } -OMATCOPY2_LAUNCHER(float, "unimplemented") -OMATCOPY2_LAUNCHER(double, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER(float) +OMATCOPY2_LAUNCHER(double) +OMATCOPY2_LAUNCHER(std::complex) +OMATCOPY2_LAUNCHER(std::complex) #undef OMATCOPY2_LAUNCHER @@ -319,25 +318,19 @@ OMATCOPY_LAUNCHER_USM(std::complex, rocblas_zgeam) #undef OMATCOPY_LAUNCHER_USM -template -sycl::event omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, - int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, - int64_t ldb, int64_t strideb, const std::vector& dependencies) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER_USM(TYPE, ROCBLAS_ROUTINE) \ - sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ - int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2(#ROCBLAS_ROUTINE, ROCBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, \ - lda, b, ldb, strideb, dependencies); \ +#define OMATCOPY2_LAUNCHER_USM(TYPE) \ + sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ + int64_t strideb, const std::vector& dependencies) { \ + return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::col_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb, \ + dependencies); \ } -OMATCOPY2_LAUNCHER_USM(float, "unimplemented") -OMATCOPY2_LAUNCHER_USM(double, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER_USM(float) +OMATCOPY2_LAUNCHER_USM(double) +OMATCOPY2_LAUNCHER_USM(std::complex) +OMATCOPY2_LAUNCHER_USM(std::complex) #undef OMATCOPY2_LAUNCHER_USM @@ -489,25 +482,19 @@ OMATCOPY_LAUNCHER(std::complex, rocblas_zgeam) #undef OMATCOPY_LAUNCHER -template -void omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, int64_t m, - int64_t n, T alpha, sycl::buffer& a, int64_t lda, std::int64_t stridea, - sycl::buffer& b, int64_t ldb, std::int64_t strideb) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER(TYPE, ROCBLAS_ROUTINE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2(#ROCBLAS_ROUTINE, ROCBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, lda, \ - b, ldb, strideb); \ +// See the column-major overloads for why geam is not used at unit stride. +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::row_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb); \ } -OMATCOPY2_LAUNCHER(float, "unimplemented") -OMATCOPY2_LAUNCHER(double, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER(float) +OMATCOPY2_LAUNCHER(double) +OMATCOPY2_LAUNCHER(std::complex) +OMATCOPY2_LAUNCHER(std::complex) #undef OMATCOPY2_LAUNCHER @@ -638,25 +625,19 @@ OMATCOPY_LAUNCHER_USM(std::complex, rocblas_zgeam) #undef OMATCOPY_LAUNCHER_USM -template -sycl::event omatcopy2(const char* func_name, Func func, sycl::queue& queue, transpose trans, - int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, - int64_t ldb, int64_t strideb, const std::vector& dependencies) { - throw unimplemented("blas", "omatcopy2", ""); -} - -#define OMATCOPY2_LAUNCHER_USM(TYPE, ROCBLAS_ROUTINE) \ - sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ - int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2(#ROCBLAS_ROUTINE, ROCBLAS_ROUTINE, queue, trans, m, n, alpha, a, stridea, \ - lda, b, ldb, strideb, dependencies); \ +#define OMATCOPY2_LAUNCHER_USM(TYPE) \ + sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ + int64_t strideb, const std::vector& dependencies) { \ + return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::row_major, trans, m, \ + n, alpha, a, lda, stridea, b, ldb, strideb, \ + dependencies); \ } -OMATCOPY2_LAUNCHER_USM(float, "unimplemented") -OMATCOPY2_LAUNCHER_USM(double, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") -OMATCOPY2_LAUNCHER_USM(std::complex, "unimplemented") +OMATCOPY2_LAUNCHER_USM(float) +OMATCOPY2_LAUNCHER_USM(double) +OMATCOPY2_LAUNCHER_USM(std::complex) +OMATCOPY2_LAUNCHER_USM(std::complex) #undef OMATCOPY2_LAUNCHER_USM From 2d018e1a915a32d076878b800e4f32aeafde3c16 Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Sat, 15 Aug 2026 09:12:09 -0700 Subject: [PATCH 2/3] [blas] Tune the omatcopy2 tile geometry for NVIDIA The tile extents were chosen on gfx950, where a wave is 64 items wide, and the strided table picks 1024-item work-groups for the 8-byte types. That is the whole of an SM's thread budget on NVIDIA and leaves only two groups resident, which measures 4% slower on double and 2% on complex across an A100 sweep, and 7% and 5% once the matrices fit in L2 and the copy is no longer bounded by main memory. Give the two backends separate strided tables rather than one compromise, as each is built against a single vendor's runtime. The unit-stride table is unchanged: the two machines agree on it to within 0.4% for every type. Every variant still fits the 17 KB budget, which on NVIDIA covers the 64 KB parts. Also record that geam was measured, not assumed, to be the slower option at unit stride on NVIDIA: on an A100 the kernels win from 1024x1024 up, by 2-231% pipelined and by up to 29% with a wait after every call. Co-authored-by: Cursor --- .../backends/cublas/cublas_extensions.cpp | 49 +++++++++-------- src/blas/backends/omatcopy2_kernels.hpp | 54 ++++++++++++++----- .../backends/rocblas/rocblas_extensions.cpp | 43 ++++++++------- 3 files changed, 94 insertions(+), 52 deletions(-) diff --git a/src/blas/backends/cublas/cublas_extensions.cpp b/src/blas/backends/cublas/cublas_extensions.cpp index fea864769..6bf6890d7 100644 --- a/src/blas/backends/cublas/cublas_extensions.cpp +++ b/src/blas/backends/cublas/cublas_extensions.cpp @@ -26,6 +26,11 @@ namespace oneapi { namespace math { namespace blas { namespace cublas { + +// This backend only ever runs on NVIDIA devices, so the omatcopy2 kernels are +// always built with the tile extents measured there. +constexpr auto omatcopy2_target = omatcopy2_kernels::vendor::nvidia; + namespace column_major { // Buffer APIs @@ -124,15 +129,16 @@ OMATCOPY_LAUNCHER(std::complex, cublasZgeam) // Unit element strides make omatcopy2 equivalent to omatcopy, but routing them // to geam is not worth it: cuBLAS has to be called from a host task and with a -// stream synchronize, whereas these kernels are plain asynchronous SYCL. The -// equivalent rocBLAS path measured 37-77% slower than the kernels on gfx950 -// even when the caller waits after every call. -#define OMATCOPY2_LAUNCHER(TYPE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::col_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb); \ +// stream synchronize, whereas these kernels are plain asynchronous SYCL. On an +// A100 the kernels beat geam everywhere from 1024x1024 up, by 2-231% when calls +// are pipelined and by up to 29% even when the caller waits after every one. +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer( \ + queue, oneapi::math::layout::col_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb); \ } OMATCOPY2_LAUNCHER(float) @@ -305,9 +311,9 @@ OMATCOPY_LAUNCHER_USM(std::complex, cublasZgeam) sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::col_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb, \ - dependencies); \ + return omatcopy2_kernels::omatcopy2_usm( \ + queue, oneapi::math::layout::col_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb, dependencies); \ } OMATCOPY2_LAUNCHER_USM(float) @@ -478,12 +484,13 @@ OMATCOPY_LAUNCHER(std::complex, cublasZgeam) #undef OMATCOPY_LAUNCHER // See the column-major overloads for why geam is not used at unit stride. -#define OMATCOPY2_LAUNCHER(TYPE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::row_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb); \ +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer( \ + queue, oneapi::math::layout::row_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb); \ } OMATCOPY2_LAUNCHER(float) @@ -656,9 +663,9 @@ OMATCOPY_LAUNCHER_USM(std::complex, cublasZgeam) sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::row_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb, \ - dependencies); \ + return omatcopy2_kernels::omatcopy2_usm( \ + queue, oneapi::math::layout::row_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb, dependencies); \ } OMATCOPY2_LAUNCHER_USM(float) diff --git a/src/blas/backends/omatcopy2_kernels.hpp b/src/blas/backends/omatcopy2_kernels.hpp index b02a68bc4..df8e44737 100644 --- a/src/blas/backends/omatcopy2_kernels.hpp +++ b/src/blas/backends/omatcopy2_kernels.hpp @@ -50,13 +50,21 @@ namespace omatcopy2_kernels { // items: the load phase puts the lane index on the row extent and the store // phase puts it on the column extent, so `cols` sets how wide the stores are. // -// Both variants below keep the tile under 17 KB for every element type, which +// Every variant below keeps the tile under 17 KB for every element type, which // leaves room for three resident groups even on CDNA1-CDNA3 (64 KB of local -// memory per compute unit); CDNA4 has 160 KB. Extents were measured on gfx950 -// over strides 1-4 and sizes from 64 to 8192 square. +// memory per compute unit) and on the smaller NVIDIA parts (64 KB on Pascal +// and Turing); CDNA4 has 160 KB and A100 has 164 KB. + +/// The two backends that include this header are each built against a single +/// vendor's runtime, so which tuning to apply is fixed at the call site rather +/// than queried from the device. +enum class vendor { nvidia, amd }; /// Unit element strides: both global accesses are contiguous, so the two /// phases want the same width and a square tile is best. +/// +/// Measured on gfx950 and on an A100, which agree to within 0.4% on every +/// element type, so one table serves both. template struct unit_stride_geometry { static constexpr int rows = sizeof(T) <= 4 ? 64 : 32; @@ -66,13 +74,33 @@ struct unit_stride_geometry { /// Real element strides: neither access is contiguous, so widening the stores /// buys nothing and a taller tile amortises the per-tile overhead instead. +/// +/// Here the two vendors disagree, so the tables are separate. Both were +/// measured over strides 1-4 and sizes from 64 to 8192 square. +template +struct strided_geometry; + +/// Measured on gfx950. template -struct strided_geometry { +struct strided_geometry { static constexpr int rows = sizeof(T) <= 8 ? 64 : 32; static constexpr int cols = 32; static constexpr int block = sizeof(T) == 8 ? 16 : 8; }; +/// Measured on an A100. The wave64 table wants 1024-item groups for the 8-byte +/// types, which is the whole of an SM's thread budget on NVIDIA and leaves only +/// two groups resident; 256 items instead is worth 4% on `double` and 2% on +/// `complex` overall, and 7% and 5% once the matrices fit in L2, where +/// the copy is not already bounded by main memory. The narrower tile for +/// 16-byte types costs 0.5% and keeps every variant inside the same 17 KB. +template +struct strided_geometry { + static constexpr int rows = 64; + static constexpr int cols = sizeof(T) <= 8 ? 32 : 16; + static constexpr int block = 4; +}; + template struct is_complex : std::false_type {}; template @@ -164,7 +192,7 @@ void launch_trans(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T al /// Picks the tile shape from the strides, which decide whether the global /// accesses are contiguous. -template +template void launch_trans_dispatch(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T alpha, bool do_conj, AccessorA a, int64_t lda, int64_t stridea, AccessorB b, int64_t ldb, int64_t strideb) { @@ -173,12 +201,12 @@ void launch_trans_dispatch(sycl::handler& cgh, int64_t logical_m, int64_t logica stridea, b, ldb, strideb); } else { - launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, lda, - stridea, b, ldb, strideb); + launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, + lda, stridea, b, ldb, strideb); } } -template +template sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, int64_t ldb, int64_t strideb, @@ -199,8 +227,8 @@ sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, return queue.submit([&](sycl::handler& cgh) { cgh.depends_on(dependencies); if (do_trans) { - launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a, lda, stridea, b, - ldb, strideb); + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a, lda, + stridea, b, ldb, strideb); } else { launch_nontrans(cgh, logical_m, logical_n, alpha, a, lda, stridea, b, ldb, strideb); @@ -208,7 +236,7 @@ sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, }); } -template +template void omatcopy2_buffer(sycl::queue& queue, oneapi::math::layout layout, oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, sycl::buffer& a, int64_t lda, int64_t stridea, sycl::buffer& b, @@ -228,8 +256,8 @@ void omatcopy2_buffer(sycl::queue& queue, oneapi::math::layout layout, // Strided writes skip elements, so the untouched ones must be preserved. auto b_acc = b.template get_access(cgh); if (do_trans) { - launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a_acc, lda, stridea, - b_acc, ldb, strideb); + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a_acc, lda, + stridea, b_acc, ldb, strideb); } else { launch_nontrans(cgh, logical_m, logical_n, alpha, a_acc, lda, stridea, b_acc, ldb, diff --git a/src/blas/backends/rocblas/rocblas_extensions.cpp b/src/blas/backends/rocblas/rocblas_extensions.cpp index a347c103b..12206d41d 100644 --- a/src/blas/backends/rocblas/rocblas_extensions.cpp +++ b/src/blas/backends/rocblas/rocblas_extensions.cpp @@ -30,6 +30,11 @@ namespace oneapi { namespace math { namespace blas { namespace rocblas { + +// This backend only ever runs on AMD devices, so the omatcopy2 kernels are +// always built with the tile extents measured there. +constexpr auto omatcopy2_target = omatcopy2_kernels::vendor::amd; + namespace column_major { // Buffer APIs @@ -136,12 +141,13 @@ OMATCOPY_LAUNCHER(std::complex, rocblas_zgeam) // stream synchronize, whereas these kernels are plain asynchronous SYCL. On // gfx950 the kernels beat the geam path by 37-77% even when the caller waits // after every call, and by more when calls can overlap. -#define OMATCOPY2_LAUNCHER(TYPE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::col_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb); \ +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer( \ + queue, oneapi::math::layout::col_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb); \ } OMATCOPY2_LAUNCHER(float) @@ -322,9 +328,9 @@ OMATCOPY_LAUNCHER_USM(std::complex, rocblas_zgeam) sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::col_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb, \ - dependencies); \ + return omatcopy2_kernels::omatcopy2_usm( \ + queue, oneapi::math::layout::col_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb, dependencies); \ } OMATCOPY2_LAUNCHER_USM(float) @@ -483,12 +489,13 @@ OMATCOPY_LAUNCHER(std::complex, rocblas_zgeam) #undef OMATCOPY_LAUNCHER // See the column-major overloads for why geam is not used at unit stride. -#define OMATCOPY2_LAUNCHER(TYPE) \ - void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ - sycl::buffer& a, int64_t lda, int64_t stridea, \ - sycl::buffer& b, int64_t ldb, int64_t strideb) { \ - omatcopy2_kernels::omatcopy2_buffer(queue, oneapi::math::layout::row_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb); \ +#define OMATCOPY2_LAUNCHER(TYPE) \ + void omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ + sycl::buffer& a, int64_t lda, int64_t stridea, \ + sycl::buffer& b, int64_t ldb, int64_t strideb) { \ + omatcopy2_kernels::omatcopy2_buffer( \ + queue, oneapi::math::layout::row_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb); \ } OMATCOPY2_LAUNCHER(float) @@ -629,9 +636,9 @@ OMATCOPY_LAUNCHER_USM(std::complex, rocblas_zgeam) sycl::event omatcopy2(sycl::queue& queue, transpose trans, int64_t m, int64_t n, TYPE alpha, \ const TYPE* a, int64_t lda, int64_t stridea, TYPE* b, int64_t ldb, \ int64_t strideb, const std::vector& dependencies) { \ - return omatcopy2_kernels::omatcopy2_usm(queue, oneapi::math::layout::row_major, trans, m, \ - n, alpha, a, lda, stridea, b, ldb, strideb, \ - dependencies); \ + return omatcopy2_kernels::omatcopy2_usm( \ + queue, oneapi::math::layout::row_major, trans, m, n, alpha, a, lda, stridea, b, ldb, \ + strideb, dependencies); \ } OMATCOPY2_LAUNCHER_USM(float) From 3fee43b056ff91d68ffb76e8b0dbff78291dc334 Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Sat, 15 Aug 2026 14:46:48 -0700 Subject: [PATCH 3/3] [blas] Tune omatcopy2 geometry for AMD architectures Select measured gfx90a and gfx942 tile geometries at runtime while retaining the gfx950 defaults for other AMD devices. Co-authored-by: Cursor --- .../backends/cublas/cublas_extensions.cpp | 2 +- src/blas/backends/omatcopy2_kernels.hpp | 136 ++++++++++++++---- .../backends/rocblas/rocblas_extensions.cpp | 6 +- 3 files changed, 115 insertions(+), 29 deletions(-) diff --git a/src/blas/backends/cublas/cublas_extensions.cpp b/src/blas/backends/cublas/cublas_extensions.cpp index 6bf6890d7..975df0d88 100644 --- a/src/blas/backends/cublas/cublas_extensions.cpp +++ b/src/blas/backends/cublas/cublas_extensions.cpp @@ -29,7 +29,7 @@ namespace cublas { // This backend only ever runs on NVIDIA devices, so the omatcopy2 kernels are // always built with the tile extents measured there. -constexpr auto omatcopy2_target = omatcopy2_kernels::vendor::nvidia; +constexpr auto omatcopy2_target = omatcopy2_kernels::target::nvidia; namespace column_major { diff --git a/src/blas/backends/omatcopy2_kernels.hpp b/src/blas/backends/omatcopy2_kernels.hpp index df8e44737..741621201 100644 --- a/src/blas/backends/omatcopy2_kernels.hpp +++ b/src/blas/backends/omatcopy2_kernels.hpp @@ -32,7 +32,9 @@ #include #include +#include #include +#include #include #include "oneapi/math/types.hpp" @@ -50,44 +52,104 @@ namespace omatcopy2_kernels { // items: the load phase puts the lane index on the row extent and the store // phase puts it on the column extent, so `cols` sets how wide the stores are. // -// Every variant below keeps the tile under 17 KB for every element type, which -// leaves room for three resident groups even on CDNA1-CDNA3 (64 KB of local -// memory per compute unit) and on the smaller NVIDIA parts (64 KB on Pascal -// and Turing); CDNA4 has 160 KB and A100 has 164 KB. +// Every variant below keeps the tile under 17 KB except the gfx90a +// complex strided tile, which uses 32.5 KB. CDNA1-CDNA3 have 64 KB of +// local memory per compute unit; CDNA4 has 160 KB and A100 has 164 KB. /// The two backends that include this header are each built against a single -/// vendor's runtime, so which tuning to apply is fixed at the call site rather -/// than queried from the device. -enum class vendor { nvidia, amd }; +/// vendor's runtime. gfx90a and gfx942 use architecture-specific tuning selected +/// at run time within the AMD backend; other AMD architectures use gfx950 +/// tuning. +enum class target { nvidia, amd, amd_gfx90a, amd_gfx942 }; + +inline target get_amd_target(const sycl::device& device) { + static thread_local std::vector> cache; + for (const auto& cached : cache) { + if (cached.first == device) { + return cached.second; + } + } + + target result = target::amd; +#ifdef SYCL_EXT_ONEAPI_DEVICE_ARCHITECTURE + // This experimental query is non-const in older DPC++ releases. + auto query_device = device; + using architecture = sycl::ext::oneapi::experimental::architecture; + if (query_device.ext_oneapi_architecture_is(architecture::amd_gpu_gfx90a)) { + result = target::amd_gfx90a; + } + else if (query_device.ext_oneapi_architecture_is(architecture::amd_gpu_gfx942)) { + result = target::amd_gfx942; + } +#endif + if (result == target::amd) { + // The DPC++ HIP adapter currently reports an unknown architecture + // through the extension above, and AdaptiveCpp does not expose it. + const auto name = device.get_info(); + if (name.find("MI210") != std::string::npos || name.find("MI250") != std::string::npos || + name.find("gfx90a") != std::string::npos) { + result = target::amd_gfx90a; + } + else if (name.find("MI300") != std::string::npos || + name.find("MI308") != std::string::npos || + name.find("MI325") != std::string::npos || + name.find("gfx942") != std::string::npos) { + result = target::amd_gfx942; + } + } + cache.emplace_back(device, result); + return result; +} /// Unit element strides: both global accesses are contiguous, so the two /// phases want the same width and a square tile is best. /// /// Measured on gfx950 and on an A100, which agree to within 0.4% on every -/// element type, so one table serves both. -template +/// element type. gfx942 retains those choices except for 16-byte elements, +/// where halving the column extent gains about 2.4%. +template struct unit_stride_geometry { static constexpr int rows = sizeof(T) <= 4 ? 64 : 32; - static constexpr int cols = rows; + static constexpr int cols = Target == target::amd_gfx942 && sizeof(T) == 16 ? 16 : rows; static constexpr int block = 8; }; /// Real element strides: neither access is contiguous, so widening the stores /// buys nothing and a taller tile amortises the per-tile overhead instead. /// -/// Here the two vendors disagree, so the tables are separate. Both were -/// measured over strides 1-4 and sizes from 64 to 8192 square. -template +/// The vendors disagree, and gfx90a and gfx942 also require targeted +/// architecture-specific choices, so these tables are separate. +template struct strided_geometry; /// Measured on gfx950. template -struct strided_geometry { +struct strided_geometry { static constexpr int rows = sizeof(T) <= 8 ? 64 : 32; static constexpr int cols = 32; static constexpr int block = sizeof(T) == 8 ? 16 : 8; }; +/// Measured on two MI210 devices. Sixteen-byte elements prefer a 64 x 32 tile +/// and a 256-item work-group, gaining about 8.5% over the gfx950 choice. Other +/// element types retain the gfx950 geometry. +template +struct strided_geometry { + static constexpr int rows = 64; + static constexpr int cols = 32; + static constexpr int block = sizeof(T) == 8 ? 16 : (sizeof(T) == 16 ? 4 : 8); +}; + +/// Measured on gfx942. Eight-byte elements prefer a 32 x 8 tile and a +/// 256-item work-group; the gfx950 choice uses 1024 items. Float and 16-byte +/// elements retain the gfx950 geometry. +template +struct strided_geometry { + static constexpr int rows = sizeof(T) <= 4 ? 64 : 32; + static constexpr int cols = sizeof(T) == 8 ? 8 : 32; + static constexpr int block = 8; +}; + /// Measured on an A100. The wave64 table wants 1024-item groups for the 8-byte /// types, which is the whole of an SM's thread budget on NVIDIA and leaves only /// two groups resident; 256 items instead is worth 4% on `double` and 2% on @@ -95,7 +157,7 @@ struct strided_geometry { /// the copy is not already bounded by main memory. The narrower tile for /// 16-byte types costs 0.5% and keeps every variant inside the same 17 KB. template -struct strided_geometry { +struct strided_geometry { static constexpr int rows = 64; static constexpr int cols = sizeof(T) <= 8 ? 32 : 16; static constexpr int block = 4; @@ -192,21 +254,41 @@ void launch_trans(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T al /// Picks the tile shape from the strides, which decide whether the global /// accesses are contiguous. -template +template void launch_trans_dispatch(sycl::handler& cgh, int64_t logical_m, int64_t logical_n, T alpha, bool do_conj, AccessorA a, int64_t lda, int64_t stridea, AccessorB b, int64_t ldb, int64_t strideb) { if (stridea == 1 && strideb == 1) { - launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, lda, - stridea, b, ldb, strideb); + launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, + a, lda, stridea, b, ldb, strideb); } else { - launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, + launch_trans>(cgh, logical_m, logical_n, alpha, do_conj, a, lda, stridea, b, ldb, strideb); } } -template +template +void launch_trans_for_device(sycl::handler& cgh, target device_target, int64_t logical_m, + int64_t logical_n, T alpha, bool do_conj, AccessorA a, int64_t lda, + int64_t stridea, AccessorB b, int64_t ldb, int64_t strideb) { + if constexpr (Target == target::amd) { + if (device_target == target::amd_gfx90a) { + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, + a, lda, stridea, b, ldb, strideb); + return; + } + if (device_target == target::amd_gfx942) { + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, + a, lda, stridea, b, ldb, strideb); + return; + } + } + launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a, lda, stridea, b, + ldb, strideb); +} + +template sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, const T* a, int64_t lda, int64_t stridea, T* b, int64_t ldb, int64_t strideb, @@ -224,11 +306,13 @@ sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, }); } + const target device_target = + do_trans && Target == target::amd ? get_amd_target(queue.get_device()) : Target; return queue.submit([&](sycl::handler& cgh) { cgh.depends_on(dependencies); if (do_trans) { - launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a, lda, - stridea, b, ldb, strideb); + launch_trans_for_device(cgh, device_target, logical_m, logical_n, alpha, + do_conj, a, lda, stridea, b, ldb, strideb); } else { launch_nontrans(cgh, logical_m, logical_n, alpha, a, lda, stridea, b, ldb, strideb); @@ -236,7 +320,7 @@ sycl::event omatcopy2_usm(sycl::queue& queue, oneapi::math::layout layout, }); } -template +template void omatcopy2_buffer(sycl::queue& queue, oneapi::math::layout layout, oneapi::math::transpose trans, int64_t m, int64_t n, T alpha, sycl::buffer& a, int64_t lda, int64_t stridea, sycl::buffer& b, @@ -251,13 +335,15 @@ void omatcopy2_buffer(sycl::queue& queue, oneapi::math::layout layout, return; } + const target device_target = + do_trans && Target == target::amd ? get_amd_target(queue.get_device()) : Target; queue.submit([&](sycl::handler& cgh) { auto a_acc = a.template get_access(cgh); // Strided writes skip elements, so the untouched ones must be preserved. auto b_acc = b.template get_access(cgh); if (do_trans) { - launch_trans_dispatch(cgh, logical_m, logical_n, alpha, do_conj, a_acc, lda, - stridea, b_acc, ldb, strideb); + launch_trans_for_device(cgh, device_target, logical_m, logical_n, alpha, + do_conj, a_acc, lda, stridea, b_acc, ldb, strideb); } else { launch_nontrans(cgh, logical_m, logical_n, alpha, a_acc, lda, stridea, b_acc, ldb, diff --git a/src/blas/backends/rocblas/rocblas_extensions.cpp b/src/blas/backends/rocblas/rocblas_extensions.cpp index 12206d41d..b5cac8bd0 100644 --- a/src/blas/backends/rocblas/rocblas_extensions.cpp +++ b/src/blas/backends/rocblas/rocblas_extensions.cpp @@ -31,9 +31,9 @@ namespace math { namespace blas { namespace rocblas { -// This backend only ever runs on AMD devices, so the omatcopy2 kernels are -// always built with the tile extents measured there. -constexpr auto omatcopy2_target = omatcopy2_kernels::vendor::amd; +// This backend only ever runs on AMD devices. The kernel dispatcher refines +// this to the gfx90a or gfx942 table at run time when needed. +constexpr auto omatcopy2_target = omatcopy2_kernels::target::amd; namespace column_major {