derek/TheRock
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
4849f1eea335bb3333f3fbf66b138f72eb85cf6c
TheRock/rocm-libraries/projects/miopen/src/gemm_v2.cpp
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
Details
CodeQL / Analyze (python) (push) Canceled after 0s
Details
pre-commit / pre-commit (push) Canceled after 0s
Details
Multi-Arch CI / setup (push) Canceled after 0s
Details
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Details
Setup (Multi-Arch) / setup (push) Canceled after 0s
Details
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Details
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
Details
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Details
Multi-Arch CI / CI Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Details
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Details
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

2431 lines
93 KiB
C++
Raw Blame History

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <miopen/config.h>
#include <miopen/gemm_v2.hpp>
#include <miopen/logger.hpp>
#include <miopen/env.hpp>
#include <miopen/tensor.hpp>
#include <miopen/handle.hpp>
#include <miopen/datatype.hpp>
#include <miopen/hipoc_kernel.hpp>
#if MIOPEN_USE_HIPBLASLT
#include <hipblaslt/hipblaslt.h>
// Only enable BF8 support if hipBLASLt is 0.8 or above
#if MIOPEN_HIPBLASLT_VERSION_FLAT >= 8000
#define ENABLE_HIPBLASLT_BF8
#endif
#endif
#if MIOPEN_USE_ROCBLAS
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunused-macros"
#define ROCBLAS_BETA_FEATURES_API 1
#pragma clang diagnostic pop
#include <half/half.hpp>
#if MIOPEN_ROCBLAS_VERSION_FLAT < 2045000
#include <rocblas.h>
#else
#include <rocblas/rocblas.h>
/// rocblas_gemm_ex3 supports F8 datatypes.
#ifdef _WIN32
#define USE_ROCBLAS_GEMM_EX3 \
((MIOPEN_ROCBLAS_VERSION_FLAT >= 3000000 && MIOPEN_ROCBLAS_VERSION_FLAT < 5000000) && \
ROCBLAS_BETA_FEATURES_API)
#else
#define USE_ROCBLAS_GEMM_EX3 \
((MIOPEN_ROCBLAS_VERSION_FLAT >= 2047000 && MIOPEN_ROCBLAS_VERSION_FLAT < 5000000) && \
ROCBLAS_BETA_FEATURES_API)
#endif
#endif
#include <miopen/perf_field.hpp>
#endif
#include <ranges>
#include <tuple> // std::ignore
#if MIOPEN_USE_ROCBLAS
/// Avoid warnings "The workspace_size and workspace arguments are obsolete" and
/// "disabled expansion of recursive macro" injected by rocblas headers.
#define AVOID_ROCBLAS_WRAPPERS_204 (MIOPEN_ROCBLAS_VERSION_FLAT >= 2004000)
/// Maintain API compatibility for versions not supporting FP16 alternate implementations
#define USE_GEMM_FLAGS_FP16_ALT_IMPL (MIOPEN_ROCBLAS_VERSION_FLAT >= 2043000)
/// Some 2.42 versions have rocblas_gemm_flags_fp16_alt_impl, but
/// some do not, and that leads to build errors.
/// Let's pass literal value as a workaround; there should be no harm.
#define USE_GEMM_FLAGS_FP16_ALT_IMPL_242 (MIOPEN_ROCBLAS_VERSION_FLAT == 2042000)
static inline uint32_t
FlagsForRocblasFp32Fp16Call(const miopen::GemmDescriptor& desc) // bool gfx90aFp16Alt)
{
#if USE_GEMM_FLAGS_FP16_ALT_IMPL
return desc.gfx90a_alt_impl ? rocblas_gemm_flags_fp16_alt_impl : 0;
#elif USE_GEMM_FLAGS_FP16_ALT_IMPL_242
return desc.gfx90a_alt_impl ? 0x4 : 0;
#else
std::ignore = desc;
MIOPEN_LOG_W("The gfx90aFp16Alt is not supported by rocBlas");
return 0;
#endif
#if USE_GEMM_FLAGS_FP16_ALT_IMPL_242 // -warning: macro is not used
#endif
}
#if USE_ROCBLAS_GEMM_EX3
static inline rocblas_computetype rocBlasComputeType_ex3(const miopen::GemmDescriptor& desc)
{
if(desc.a_cast_type == miopenFloat8_fnuz && desc.b_cast_type == miopenFloat8_fnuz)
{
return rocblas_compute_type_f8_f8_f32;
}
else if(desc.a_cast_type == miopenFloat8_fnuz && desc.b_cast_type == miopenBFloat8_fnuz)
{
return rocblas_compute_type_f8_bf8_f32;
}
else if(desc.a_cast_type == miopenBFloat8_fnuz && desc.b_cast_type == miopenFloat8_fnuz)
{
return rocblas_compute_type_bf8_f8_f32;
}
else if(desc.a_cast_type == miopenBFloat8_fnuz && desc.b_cast_type == miopenBFloat8_fnuz)
{
return rocblas_compute_type_bf8_bf8_f32;
}
else
{
return rocblas_compute_type_f32;
}
}
#endif
static inline rocblas_datatype rocBlasComputeType(const miopen::GemmDescriptor& desc)
{
if(desc.dataType == miopenInt8)
return rocblas_datatype::rocblas_datatype_i32_r;
else
return rocblas_datatype::rocblas_datatype_f32_r;
}
auto rocBlasDataType(miopenDataType_t data_type)
{
/// \todo Not all supported data types are handled here.
/// This is fine so far because this function is used only with FP16/F8.
#if USE_ROCBLAS_GEMM_EX3
if(data_type == miopenFloat8_fnuz)
return rocblas_datatype::rocblas_datatype_f8_r;
if(data_type == miopenBFloat8_fnuz)
return rocblas_datatype::rocblas_datatype_bf8_r;
#endif
if(data_type == miopenHalf)
return rocblas_datatype::rocblas_datatype_f16_r;
MIOPEN_THROW(miopenStatusInternalError, "Invalid data type passed");
}
template <typename T>
rocblas_status miopen_rocblas_gemm_ex3(const miopen::Handle& handle,
const miopen::GemmDescriptor& gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset)
{
rocblas_status rb_status =
rocblas_status::rocblas_status_internal_error; // cppcheck-suppress redundantInitialization
#if USE_ROCBLAS_GEMM_EX3
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdisabled-macro-expansion"
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
auto flags = FlagsForRocblasFp32Fp16Call(gemm_desc);
if(gemm_desc.conv_attributes.fp8rounding_mode.Get() == miopenF8RoundingModeStochastic)
flags = flags | rocblas_gemm_flags::rocblas_gemm_flags_stochastic_rounding;
rb_status = // cppcheck-suppress redundantInitialization
rocblas_gemm_ex3(handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const T*>(A) + a_offset,
rocBlasDataType(gemm_desc.dataType),
gemm_desc.lda,
static_cast<const T*>(B) + b_offset,
rocBlasDataType(gemm_desc.dataType),
gemm_desc.ldb,
&beta,
static_cast<const T*>(C) + c_offset,
rocBlasDataType(gemm_desc.dataType),
gemm_desc.ldc,
static_cast<T*>(C) + c_offset,
rocBlasDataType(gemm_desc.dataType),
gemm_desc.ldc,
rocBlasComputeType_ex3(gemm_desc),
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
flags); // gfx90a_alt_impl));
return rb_status;
#pragma clang diagnostic pop
#else
std::ignore = A;
std::ignore = a_offset;
std::ignore = B;
std::ignore = b_offset;
std::ignore = C;
std::ignore = c_offset;
#endif
MIOPEN_THROW(miopenStatusInternalError,
"An appropriate version of rocBLAS is required for this op");
std::ignore = handle;
std::ignore = gemm_desc;
return rb_status;
}
template <class... Ts>
auto miopen_rocblas_gemm_ex(const miopen::Handle& handle,
const miopen::GemmDescriptor& gemm_desc,
Ts... xs)
{
std::ignore = handle;
std::ignore = gemm_desc;
#if AVOID_ROCBLAS_WRAPPERS_204
return (rocblas_gemm_ex)(handle.rhandle().get(), xs...);
#else
std::size_t zero = 0;
return rocblas_gemm_ex(handle.rhandle().get(), xs..., &zero, nullptr);
#endif
}
template <class... Ts>
auto miopen_rocblas_gemm_strided_batched_ex(Ts... xs)
{
#if AVOID_ROCBLAS_WRAPPERS_204
return (rocblas_gemm_strided_batched_ex)(xs...);
#else
std::size_t zero = 0;
return rocblas_gemm_strided_batched_ex(xs..., &zero, nullptr);
#endif
}
template <typename T>
rocblas_status miopen_rocblas_gemm_strided_batched_ex3(const miopen::Handle& handle,
const miopen::GemmDescriptor& gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset)
{
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
// Until there is a batched counter part to the ex3 rocBlas call we need to iterate over the
// batched GEMM
for(int bCount = 0; bCount < gemm_desc.batch_count; ++bCount)
{
rb_status = miopen_rocblas_gemm_ex3<T>(handle,
gemm_desc,
A,
a_offset + (bCount * gemm_desc.strideA),
B,
b_offset + (bCount * gemm_desc.strideB),
C,
c_offset + (bCount * gemm_desc.strideC));
}
return rb_status;
}
#endif // MIOPEN_USE_ROCBLAS
MIOPEN_DECLARE_ENV_VAR_UINT64(MIOPEN_GEMM_ENFORCE_BACKEND)
namespace miopen {
bool IsFP8Supported(const std::string& device_name)
{
#if defined(USE_ROCBLAS_GEMM_EX3) && USE_ROCBLAS_GEMM_EX3
return device_name == "gfx942" || miopen::StartsWith(device_name, "gfx95");
#else
std::ignore = device_name;
return false;
#endif
}
std::ostream& operator<<(std::ostream& stream, const GemmDescriptor& gemm_desc)
{
return stream << "{" << "isColMajor " << gemm_desc.isColMajor << ", " << "transA "
<< gemm_desc.transA << ", " << "transB " << gemm_desc.transB << ", " << "m "
<< gemm_desc.m << ", " << "n " << gemm_desc.n << ", " << "k " << gemm_desc.k
<< ", " << "lda " << gemm_desc.lda << ", " << "ldb " << gemm_desc.ldb << ", "
<< "ldc " << gemm_desc.ldc << ", " << "batch_count " << gemm_desc.batch_count
<< ", " << "strideA " << gemm_desc.strideA << ", " << "strideB "
<< gemm_desc.strideB << ", " << "strideC " << gemm_desc.strideC << ", "
<< "alpha " << gemm_desc.alpha << ", " << "beta " << gemm_desc.beta << ", "
<< "dataType " << GetDataType(gemm_desc.dataType) << ", " << "a_cast_type "
<< GetDataType(gemm_desc.a_cast_type) << ", " << "b_cast_type "
<< GetDataType(gemm_desc.b_cast_type) << "} ";
}
#if MIOPEN_USE_ROCBLAS
inline rocblas_atomics_mode DisableRocblasAtomics(const miopen::Handle& handle)
{
MIOPEN_LOG_I2("");
rocblas_atomics_mode cur_mode;
[[maybe_unused]] rocblas_status status =
rocblas_get_atomics_mode(handle.rhandle().get(), &cur_mode);
assert(status == rocblas_status::rocblas_status_success);
if(cur_mode == rocblas_atomics_allowed)
{
status = rocblas_set_atomics_mode(handle.rhandle().get(), rocblas_atomics_not_allowed);
assert(status == rocblas_status::rocblas_status_success);
}
return cur_mode;
}
inline void SetRocblasAtomics(const miopen::Handle& handle, rocblas_atomics_mode mode)
{
MIOPEN_LOG_I2("");
[[maybe_unused]] rocblas_status status = rocblas_set_atomics_mode(handle.rhandle().get(), mode);
assert(status == rocblas_status::rocblas_status_success);
}
#endif
inline void ProfilingRecordStart(const Handle& handle, HipEventPtr& start, HipEventPtr& stop)
{
start = make_hip_event();
stop = make_hip_event();
(void)hipEventRecord(start.get(), handle.GetStream());
}
inline void ProfilingRecordStop(const Handle& handle, HipEventPtr& start, HipEventPtr& stop)
{
(void)hipEventRecord(stop.get(), handle.GetStream());
(void)hipEventSynchronize(stop.get());
float mS = 0;
(void)hipEventElapsedTime(&mS, start.get(), stop.get());
handle.ResetKernelTime();
handle.AccumKernelTime(mS);
}
#if MIOPEN_USE_HIPBLASLT
struct HipBLASLtMemoryHandles
{
hipblasLtMatrixLayout_t matA, matB, matC, matD;
hipblasLtMatmulDesc_t matmul;
hipblasLtMatmulPreference_t pref;
HipBLASLtMemoryHandles()
: matA(nullptr), matB(nullptr), matC(nullptr), matD(nullptr), matmul(nullptr), pref(nullptr)
{
}
~HipBLASLtMemoryHandles()
{
hipblasLtMatrixLayoutDestroy(matA);
hipblasLtMatrixLayoutDestroy(matB);
hipblasLtMatrixLayoutDestroy(matC);
hipblasLtMatrixLayoutDestroy(matD);
hipblasLtMatmulDescDestroy(matmul);
hipblasLtMatmulPreferenceDestroy(pref);
}
};
static inline void check_hipblas_status(hipblasStatus_t status)
{
if(status != hipblasStatus_t::HIPBLAS_STATUS_SUCCESS)
{
std::cout << "error msg" << std::endl;
MIOPEN_THROW(miopenStatusInternalError, "hipBlasLt error encountered");
}
}
template <typename DataTypeAB, typename DataTypeC>
static void miopen_hipblasLt_gemm(const miopen::Handle& handle,
const miopen::GemmDescriptor& gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
hipDataType hip_type_AB,
Data_t C,
std::size_t c_offset,
hipDataType hip_type_C,
bool skip_batches)
{
HipBLASLtMemoryHandles hipBLASLtHandles;
if(gemm_desc.transA)
{
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matA, hip_type_AB, gemm_desc.k, gemm_desc.m, gemm_desc.lda));
}
else
{
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matA, hip_type_AB, gemm_desc.m, gemm_desc.k, gemm_desc.lda));
}
if(gemm_desc.transB)
{
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matB, hip_type_AB, gemm_desc.n, gemm_desc.k, gemm_desc.ldb));
}
else
{
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matB, hip_type_AB, gemm_desc.k, gemm_desc.n, gemm_desc.ldb));
}
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matC, hip_type_C, gemm_desc.m, gemm_desc.n, gemm_desc.ldc));
check_hipblas_status(hipblasLtMatrixLayoutCreate(
&hipBLASLtHandles.matD, hip_type_C, gemm_desc.m, gemm_desc.n, gemm_desc.ldc));
if(gemm_desc.batch_count > 1 && !skip_batches)
{
check_hipblas_status(hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matA,
HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&gemm_desc.batch_count,
sizeof(gemm_desc.batch_count)));
check_hipblas_status(
hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matA,
HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&gemm_desc.strideA,
sizeof(gemm_desc.strideA)));
check_hipblas_status(hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matB,
HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&gemm_desc.batch_count,
sizeof(gemm_desc.batch_count)));
check_hipblas_status(
hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matB,
HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&gemm_desc.strideB,
sizeof(gemm_desc.strideB)));
check_hipblas_status(hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matC,
HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&gemm_desc.batch_count,
sizeof(gemm_desc.batch_count)));
check_hipblas_status(
hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matC,
HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&gemm_desc.strideC,
sizeof(gemm_desc.strideC)));
check_hipblas_status(hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matD,
HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&gemm_desc.batch_count,
sizeof(gemm_desc.batch_count)));
check_hipblas_status(
hipblasLtMatrixLayoutSetAttribute(hipBLASLtHandles.matD,
HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&gemm_desc.strideC,
sizeof(gemm_desc.strideC)));
}
check_hipblas_status(
hipblasLtMatmulDescCreate(&hipBLASLtHandles.matmul, HIPBLAS_COMPUTE_32F, HIP_R_32F));
hipblasOperation_t opTypeA = gemm_desc.transA ? HIPBLAS_OP_T : HIPBLAS_OP_N;
hipblasOperation_t opTypeB = gemm_desc.transB ? HIPBLAS_OP_T : HIPBLAS_OP_N;
check_hipblas_status(hipblasLtMatmulDescSetAttribute(
hipBLASLtHandles.matmul, HIPBLASLT_MATMUL_DESC_TRANSA, &opTypeA, sizeof(opTypeA)));
check_hipblas_status(hipblasLtMatmulDescSetAttribute(
hipBLASLtHandles.matmul, HIPBLASLT_MATMUL_DESC_TRANSB, &opTypeB, sizeof(opTypeB)));
hipblasLtEpilogue_t epilogue = HIPBLASLT_EPILOGUE_DEFAULT;
check_hipblas_status(hipblasLtMatmulDescSetAttribute(
hipBLASLtHandles.matmul, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
/// \todo Need to request additional workspace for optimal gemm performance, and pass down
/// workspace size & pointer. --BrianHarrisonAMD June 2024
size_t max_workspace_size = 0;
void* workspace = nullptr;
check_hipblas_status(hipblasLtMatmulPreferenceCreate(&hipBLASLtHandles.pref));
check_hipblas_status(
hipblasLtMatmulPreferenceSetAttribute(hipBLASLtHandles.pref,
HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&max_workspace_size,
sizeof(max_workspace_size)));
const int requestSolutions = 1;
hipblasLtMatmulHeuristicResult_t heuristicResult[requestSolutions];
int returnedAlgoCount = 0;
check_hipblas_status(hipblasLtMatmulAlgoGetHeuristic(handle.HipblasLtHandle().get(),
hipBLASLtHandles.matmul,
hipBLASLtHandles.matA,
hipBLASLtHandles.matB,
hipBLASLtHandles.matC,
hipBLASLtHandles.matD,
hipBLASLtHandles.pref,
requestSolutions,
heuristicResult,
&returnedAlgoCount));
if(returnedAlgoCount == 0)
{
MIOPEN_THROW(miopenStatusInternalError,
"no solution found for hipBLASLt hipBLASLtHandles.matmul");
}
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
const void* aData = static_cast<const DataTypeAB*>(A) + a_offset;
const void* bData = static_cast<const DataTypeAB*>(B) + b_offset;
const void* cData = static_cast<const DataTypeC*>(C) + c_offset;
void* dData = static_cast<DataTypeC*>(C) + c_offset;
{
HipEventProfiler profiler(handle);
check_hipblas_status(hipblasLtMatmul(handle.HipblasLtHandle().get(),
hipBLASLtHandles.matmul,
&alpha,
aData,
hipBLASLtHandles.matA,
bData,
hipBLASLtHandles.matB,
&beta,
cData,
hipBLASLtHandles.matC,
dData,
hipBLASLtHandles.matD,
&heuristicResult[0].algo,
workspace,
max_workspace_size,
handle.GetStream()));
}
}
static void call_miopen_hipblasLt_gemm(const miopen::Handle& handle,
const miopen::GemmDescriptor& gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
bool skip_batches)
{
switch(gemm_desc.dataType)
{
case miopenInt8: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt8 is not supported for hipBLASLt");
}
break;
case miopenInt32: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt32 is not supported for hipBLASLt");
}
break;
case miopenHalf: {
miopen_hipblasLt_gemm<hipblasLtHalf, hipblasLtHalf>(handle,
gemm_desc,
A,
a_offset,
B,
b_offset,
HIP_R_16F,
C,
c_offset,
HIP_R_16F,
skip_batches);
}
break;
case miopenBFloat16: {
miopen_hipblasLt_gemm<hipblasLtBfloat16, hipblasLtBfloat16>(handle,
gemm_desc,
A,
a_offset,
B,
b_offset,
HIP_R_16BF,
C,
c_offset,
HIP_R_16BF,
skip_batches);
}
break;
case miopenFloat: {
miopen_hipblasLt_gemm<hipblasLtFloat, hipblasLtFloat>(handle,
gemm_desc,
A,
a_offset,
B,
b_offset,
HIP_R_32F,
C,
c_offset,
HIP_R_32F,
skip_batches);
}
break;
case miopenFloat8_fnuz: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
if(is_gfx942)
{
miopen_hipblasLt_gemm<hipblaslt_f8_fnuz, hipblaslt_f8_fnuz>(handle,
gemm_desc,
A,
a_offset,
B,
b_offset,
HIP_R_8F_E4M3_FNUZ,
C,
c_offset,
HIP_R_8F_E4M3_FNUZ,
skip_batches);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"miopenFloat8_fnuz is only supported for hipBlasLt on gfx942");
}
}
break;
case miopenBFloat8_fnuz: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
if(is_gfx942)
{
#ifdef ENABLE_HIPBLASLT_BF8
miopen_hipblasLt_gemm<hipblaslt_bf8_fnuz, hipblaslt_bf8_fnuz>(handle,
gemm_desc,
A,
a_offset,
B,
b_offset,
HIP_R_8F_E5M2_FNUZ,
C,
c_offset,
HIP_R_8F_E5M2_FNUZ,
skip_batches);
#else
MIOPEN_THROW(
miopenStatusInternalError,
"miopenBFloat8_fnuz is not supported for this version of hipBlasLt on gfx942");
#endif
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"miopenBFloat8_fnuz is only supported for hipBlasLt on gfx942");
}
}
break;
case miopenDouble: {
MIOPEN_THROW(miopenStatusInternalError, "miopenDouble is not supported for hipBlasLt");
}
break;
case miopenInt64: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt64 is not supported for hipBlasLt");
}
break;
}
}
#endif
// hacks: control GEMM backend by enviroment variable and build option
// very nasty
static GemmBackend_t enforce_gemm_backend(GemmBackend_t gemm_backend_preferred)
{
GemmBackend_t gemm_backend_env = GemmBackend_t::nogemmbackend;
// enforce backend based on env variable
// I have left the commented lines here to preserve values for the enforce and hint at why are
// they 1, 3, and 5
switch(env::value(MIOPEN_GEMM_ENFORCE_BACKEND))
{
#if MIOPEN_USE_ROCBLAS
case 1: gemm_backend_env = GemmBackend_t::rocblas; break;
#endif
// case 2: gemm_backend_env = GemmBackend_t::miopengemm; break;
case 3:
gemm_backend_env = GemmBackend_t::nogemmbackend;
break;
// case 4: gemm_backend_env = GemmBackend_t::miopentensile; break;
#if MIOPEN_USE_HIPBLASLT
case 5: gemm_backend_env = GemmBackend_t::hipblaslt; break;
#endif
default: gemm_backend_env = gemm_backend_preferred;
}
return gemm_backend_env;
}
miopenStatus_t CallGemm(const Handle& handle,
GemmDescriptor gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
GemmBackend_t gemm_backend)
{
MIOPEN_LOG_I2("gemm_desc: " << gemm_desc);
gemm_backend = enforce_gemm_backend(gemm_backend);
if(!gemm_desc.isColMajor)
{
gemm_desc.isColMajor = !gemm_desc.isColMajor;
std::swap(A, B);
std::swap(a_offset, b_offset);
std::swap(gemm_desc.a_cast_type, gemm_desc.b_cast_type);
std::swap(gemm_desc.transA, gemm_desc.transB);
std::swap(gemm_desc.m, gemm_desc.n);
std::swap(gemm_desc.lda, gemm_desc.ldb);
}
switch(gemm_backend)
{
case GemmBackend_t::nogemmbackend: return miopenStatusNotImplemented;
case GemmBackend_t::rocblas: {
#if MIOPEN_USE_ROCBLAS
MIOPEN_LOG_I2("rocBLAS");
HipEventPtr start = nullptr;
HipEventPtr stop = nullptr;
if(handle.IsProfilingEnabled())
{
ProfilingRecordStart(handle, start, stop);
}
rocblas_atomics_mode cur_mode =
rocblas_atomics_mode::rocblas_atomics_allowed; // default value from rocblas
if(gemm_desc.deterministic)
cur_mode = DisableRocblasAtomics(handle);
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
switch(gemm_desc.dataType)
{
case miopenInt8: {
assert(gemm_desc.k % 4 == 0);
auto alpha = int(gemm_desc.alpha);
auto beta = int(gemm_desc.beta);
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const int8_t*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.lda,
static_cast<const int8_t*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_int*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
static_cast<rocblas_int*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_i32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenInt32: break;
case miopenHalf: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
// We need ex3 API if any of the dataType or the cast type is an 8-bit floating type
const auto needs_ex3 = [&]() {
if((gemm_desc.dataType == miopenFloat8_fnuz ||
gemm_desc.dataType == miopenBFloat8_fnuz) ||
(gemm_desc.a_cast_type == miopenFloat8_fnuz ||
gemm_desc.a_cast_type == miopenBFloat8_fnuz) ||
(gemm_desc.b_cast_type == miopenBFloat8_fnuz ||
gemm_desc.b_cast_type == miopenFloat8_fnuz))
{
return true;
}
else
{
return false;
}
}();
// ex3 API only works on the gfx942 ASIC;
if(needs_ex3)
{
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_ex3<rocblas_half>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
}
else
{
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_half*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.lda,
static_cast<const rocblas_half*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_half*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
static_cast<rocblas_half*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc),
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
FlagsForRocblasFp32Fp16Call(gemm_desc)); // gfx90a_alt_impl));
}
}
break;
case miopenBFloat16: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_bfloat16*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.lda,
static_cast<const rocblas_bfloat16*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_bfloat16*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
static_cast<rocblas_bfloat16*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc),
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenFloat: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const float*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.lda,
static_cast<const float*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldb,
&beta,
static_cast<const float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
static_cast<float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenFloat8_fnuz:
case miopenBFloat8_fnuz: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_ex3<char>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
};
break;
case miopenDouble: {
MIOPEN_THROW(miopenStatusInternalError,
"miopenDouble data type not supported by rocBLAS.");
};
break;
case miopenInt64: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt64 is not currently supported.");
}
break;
}
if(handle.IsProfilingEnabled())
ProfilingRecordStop(handle, start, stop);
if(rb_status != rocblas_status::rocblas_status_success)
MIOPEN_THROW(miopenStatusInternalError, "rocBlas error encountered");
if(gemm_desc.deterministic)
SetRocblasAtomics(handle, cur_mode);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
case GemmBackend_t::hipblaslt: {
#if MIOPEN_USE_HIPBLASLT
call_miopen_hipblasLt_gemm(handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset, true);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
}
return miopenStatusUnknownError;
}
miopenStatus_t CallGemm(const Handle& handle,
GemmDescriptor gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
miopenDataType_t cType,
GemmBackend_t gemm_backend)
{
// If C/D type matches A/B type, delegate to the standard overload
if(cType == gemm_desc.dataType)
return CallGemm(handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset, gemm_backend);
MIOPEN_LOG_I2("gemm_desc: " << gemm_desc << " cType: " << GetDataType(cType));
gemm_backend = enforce_gemm_backend(gemm_backend);
if(!gemm_desc.isColMajor)
{
gemm_desc.isColMajor = !gemm_desc.isColMajor;
std::swap(A, B);
std::swap(a_offset, b_offset);
std::swap(gemm_desc.a_cast_type, gemm_desc.b_cast_type);
std::swap(gemm_desc.transA, gemm_desc.transB);
std::swap(gemm_desc.m, gemm_desc.n);
std::swap(gemm_desc.lda, gemm_desc.ldb);
}
switch(gemm_backend)
{
case GemmBackend_t::nogemmbackend: return miopenStatusNotImplemented;
case GemmBackend_t::rocblas: {
#if MIOPEN_USE_ROCBLAS
MIOPEN_LOG_I2("rocBLAS mixed-precision C/D");
HipEventPtr start = nullptr;
HipEventPtr stop = nullptr;
if(handle.IsProfilingEnabled())
{
ProfilingRecordStart(handle, start, stop);
}
rocblas_atomics_mode cur_mode = rocblas_atomics_mode::rocblas_atomics_allowed;
if(gemm_desc.deterministic)
cur_mode = DisableRocblasAtomics(handle);
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
// Currently only bf16 A/B with fp32 C/D is supported
if(gemm_desc.dataType == miopenBFloat16 && cType == miopenFloat)
{
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_bfloat16*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.lda,
static_cast<const rocblas_bfloat16*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldb,
&beta,
static_cast<const float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
static_cast<float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc),
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"CallGemm with cType: unsupported dataType/cType combination");
}
if(handle.IsProfilingEnabled())
ProfilingRecordStop(handle, start, stop);
if(rb_status != rocblas_status::rocblas_status_success)
MIOPEN_THROW(miopenStatusInternalError, "rocBlas error encountered");
if(gemm_desc.deterministic)
SetRocblasAtomics(handle, cur_mode);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
case GemmBackend_t::hipblaslt: {
MIOPEN_THROW(miopenStatusInternalError,
"CallGemm with cType: hipblaslt backend not supported");
}
}
return miopenStatusUnknownError;
}
miopenStatus_t CallGemmStridedBatched(const Handle& handle,
GemmDescriptor gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
GemmBackend_t gemm_backend)
{
MIOPEN_LOG_I2("gemm_desc: " << gemm_desc);
gemm_backend = enforce_gemm_backend(gemm_backend);
if(!gemm_desc.isColMajor)
{
gemm_desc.isColMajor = !gemm_desc.isColMajor;
std::swap(A, B);
std::swap(a_offset, b_offset);
std::swap(gemm_desc.a_cast_type, gemm_desc.b_cast_type);
std::swap(gemm_desc.transA, gemm_desc.transB);
std::swap(gemm_desc.m, gemm_desc.n);
std::swap(gemm_desc.lda, gemm_desc.ldb);
std::swap(gemm_desc.strideA, gemm_desc.strideB);
}
switch(gemm_backend)
{
case GemmBackend_t::nogemmbackend: return miopenStatusNotImplemented;
case GemmBackend_t::rocblas: {
#if MIOPEN_USE_ROCBLAS
MIOPEN_LOG_I2("rocBLAS");
HipEventPtr start = nullptr;
HipEventPtr stop = nullptr;
if(handle.IsProfilingEnabled())
{
ProfilingRecordStart(handle, start, stop);
}
rocblas_atomics_mode cur_mode =
rocblas_atomics_mode::rocblas_atomics_allowed; // default value from rocblas
if(gemm_desc.deterministic)
cur_mode = DisableRocblasAtomics(handle);
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
switch(gemm_desc.dataType)
{
case miopenInt8: {
assert(gemm_desc.k % 4 == 0);
auto alpha = int(gemm_desc.alpha);
auto beta = int(gemm_desc.beta);
rb_status = miopen_rocblas_gemm_strided_batched_ex(
handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const int8_t*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.lda,
gemm_desc.strideA,
static_cast<const int8_t*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.ldb,
gemm_desc.strideB,
&beta,
static_cast<const rocblas_int*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
gemm_desc.strideC,
static_cast<rocblas_int*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
gemm_desc.strideC,
gemm_desc.batch_count,
rocblas_datatype::rocblas_datatype_i32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenInt32: break;
case miopenHalf: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
// We need ex3 API if any of the dataType or the cast type is an 8-bit floating type
const auto needs_ex3 = [&]() {
if((gemm_desc.dataType == miopenFloat8_fnuz ||
gemm_desc.dataType == miopenBFloat8_fnuz) ||
(gemm_desc.a_cast_type == miopenFloat8_fnuz ||
gemm_desc.a_cast_type == miopenBFloat8_fnuz) ||
(gemm_desc.b_cast_type == miopenBFloat8_fnuz ||
gemm_desc.b_cast_type == miopenFloat8_fnuz))
{
return true;
}
else
{
return false;
}
}();
// ex3 API only works on the gfx942 ASIC;
if(needs_ex3)
{
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_strided_batched_ex3<rocblas_half>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
}
else
{
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_strided_batched_ex(
handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_half*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.lda,
gemm_desc.strideA,
static_cast<const rocblas_half*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldb,
gemm_desc.strideB,
&beta,
static_cast<const rocblas_half*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
gemm_desc.strideC,
static_cast<rocblas_half*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
gemm_desc.strideC,
gemm_desc.batch_count,
rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
FlagsForRocblasFp32Fp16Call(gemm_desc));
}
}
break;
case miopenBFloat16: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_strided_batched_ex(
handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_bfloat16*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.lda,
gemm_desc.strideA,
static_cast<const rocblas_bfloat16*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldb,
gemm_desc.strideB,
&beta,
static_cast<const rocblas_bfloat16*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
gemm_desc.strideC,
static_cast<rocblas_bfloat16*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
gemm_desc.strideC,
gemm_desc.batch_count,
rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenFloat: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_strided_batched_ex(
handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const float*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.lda,
gemm_desc.strideA,
static_cast<const float*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldb,
gemm_desc.strideB,
&beta,
static_cast<const float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
gemm_desc.strideC,
static_cast<float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
gemm_desc.strideC,
gemm_desc.batch_count,
rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
break;
case miopenFloat8_fnuz:
case miopenBFloat8_fnuz: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_strided_batched_ex3<char>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
break;
}
case miopenDouble: {
MIOPEN_THROW(miopenStatusInternalError,
"miopenDouble data type not supported by rocBLAS.");
}
break;
case miopenInt64: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt64 is not currently supported.");
}
break;
}
if(handle.IsProfilingEnabled())
ProfilingRecordStop(handle, start, stop);
if(rb_status != rocblas_status::rocblas_status_success)
MIOPEN_THROW(miopenStatusInternalError, "rocBlas error encountered");
if(gemm_desc.deterministic)
SetRocblasAtomics(handle, cur_mode);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
case GemmBackend_t::hipblaslt: {
#if MIOPEN_USE_HIPBLASLT
call_miopen_hipblasLt_gemm(handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset, false);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
}
return miopenStatusUnknownError;
}
miopenStatus_t CallGemmStridedBatched(const Handle& handle,
GemmDescriptor gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
miopenDataType_t cType,
GemmBackend_t gemm_backend)
{
// If C/D type matches A/B type, delegate to the standard overload
if(cType == gemm_desc.dataType)
return CallGemmStridedBatched(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset, gemm_backend);
MIOPEN_LOG_I2("gemm_desc: " << gemm_desc << " cType: " << GetDataType(cType));
gemm_backend = enforce_gemm_backend(gemm_backend);
if(!gemm_desc.isColMajor)
{
gemm_desc.isColMajor = !gemm_desc.isColMajor;
std::swap(A, B);
std::swap(a_offset, b_offset);
std::swap(gemm_desc.a_cast_type, gemm_desc.b_cast_type);
std::swap(gemm_desc.transA, gemm_desc.transB);
std::swap(gemm_desc.m, gemm_desc.n);
std::swap(gemm_desc.lda, gemm_desc.ldb);
std::swap(gemm_desc.strideA, gemm_desc.strideB);
}
switch(gemm_backend)
{
case GemmBackend_t::nogemmbackend: return miopenStatusNotImplemented;
case GemmBackend_t::rocblas: {
#if MIOPEN_USE_ROCBLAS
MIOPEN_LOG_I2("rocBLAS mixed-precision C/D (strided batched)");
HipEventPtr start = nullptr;
HipEventPtr stop = nullptr;
if(handle.IsProfilingEnabled())
{
ProfilingRecordStart(handle, start, stop);
}
rocblas_atomics_mode cur_mode = rocblas_atomics_mode::rocblas_atomics_allowed;
if(gemm_desc.deterministic)
cur_mode = DisableRocblasAtomics(handle);
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
// Currently only bf16 A/B with fp32 C/D is supported
if(gemm_desc.dataType == miopenBFloat16 && cType == miopenFloat)
{
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
rb_status = miopen_rocblas_gemm_strided_batched_ex(
handle.rhandle().get(),
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_bfloat16*>(A) + a_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.lda,
gemm_desc.strideA,
static_cast<const rocblas_bfloat16*>(B) + b_offset,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldb,
gemm_desc.strideB,
&beta,
static_cast<const float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
gemm_desc.strideC,
static_cast<float*>(C) + c_offset,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
gemm_desc.strideC,
gemm_desc.batch_count,
rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
else
{
MIOPEN_THROW(
miopenStatusInternalError,
"CallGemmStridedBatched with cType: unsupported dataType/cType combination");
}
if(handle.IsProfilingEnabled())
ProfilingRecordStop(handle, start, stop);
if(rb_status != rocblas_status::rocblas_status_success)
MIOPEN_THROW(miopenStatusInternalError, "rocBlas error encountered");
if(gemm_desc.deterministic)
SetRocblasAtomics(handle, cur_mode);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
case GemmBackend_t::hipblaslt: {
MIOPEN_THROW(miopenStatusInternalError,
"CallGemmStridedBatched with cType: hipblaslt backend not supported");
}
}
return miopenStatusUnknownError;
}
miopenStatus_t CallGemmStridedBatchedSequential(const Handle& handle,
GemmDescriptor gemm_desc,
ConstData_t A,
std::size_t a_offset,
ConstData_t B,
std::size_t b_offset,
Data_t C,
std::size_t c_offset,
GemmBackend_t gemm_backend)
{
MIOPEN_LOG_I2("gemm_desc: " << gemm_desc);
gemm_backend = enforce_gemm_backend(gemm_backend);
if(!gemm_desc.isColMajor)
{
gemm_desc.isColMajor = !gemm_desc.isColMajor;
std::swap(A, B);
std::swap(a_offset, b_offset);
std::swap(gemm_desc.a_cast_type, gemm_desc.b_cast_type);
std::swap(gemm_desc.transA, gemm_desc.transB);
std::swap(gemm_desc.m, gemm_desc.n);
std::swap(gemm_desc.lda, gemm_desc.ldb);
std::swap(gemm_desc.strideA, gemm_desc.strideB);
}
switch(gemm_backend)
{
case GemmBackend_t::nogemmbackend: return miopenStatusNotImplemented;
case GemmBackend_t::rocblas: {
#if MIOPEN_USE_ROCBLAS
MIOPEN_LOG_I2("rocBLAS");
HipEventPtr start = nullptr;
HipEventPtr stop = nullptr;
if(handle.IsProfilingEnabled())
{
ProfilingRecordStart(handle, start, stop);
}
rocblas_atomics_mode cur_mode =
rocblas_atomics_mode::rocblas_atomics_allowed; // default value from rocblas
if(gemm_desc.deterministic)
{
cur_mode = DisableRocblasAtomics(handle);
}
rocblas_status rb_status = rocblas_status::rocblas_status_internal_error;
switch(gemm_desc.dataType)
{
case miopenInt8: {
assert(gemm_desc.k % 4 == 0);
auto alpha = int(gemm_desc.alpha);
auto beta = int(gemm_desc.beta);
for(int i = 0; i < gemm_desc.batch_count; ++i)
{
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const int8_t*>(A) + a_offset + i * gemm_desc.strideA,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.lda,
static_cast<const int8_t*>(B) + b_offset + i * gemm_desc.strideB,
rocblas_datatype::rocblas_datatype_i8_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_int*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
static_cast<rocblas_int*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_i32_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_i32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
}
break;
case miopenInt32: break;
case miopenHalf: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
// We need ex3 API if any of the dataType or the cast type is an 8-bit floating type
const auto needs_ex3 = [&]() {
if((gemm_desc.dataType == miopenFloat8_fnuz ||
gemm_desc.dataType == miopenBFloat8_fnuz) ||
(gemm_desc.a_cast_type == miopenFloat8_fnuz ||
gemm_desc.a_cast_type == miopenBFloat8_fnuz) ||
(gemm_desc.b_cast_type == miopenBFloat8_fnuz ||
gemm_desc.b_cast_type == miopenFloat8_fnuz))
{
return true;
}
else
{
return false;
}
}();
// ex3 API only works on the gfx942 ASIC;
if(needs_ex3)
{
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_strided_batched_ex3<rocblas_half>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
}
else
{
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
for(int i = 0; i < gemm_desc.batch_count; ++i)
{
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_half*>(A) + a_offset + i * gemm_desc.strideA,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.lda,
static_cast<const rocblas_half*>(B) + b_offset + i * gemm_desc.strideB,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_half*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
static_cast<rocblas_half*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_f16_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
FlagsForRocblasFp32Fp16Call(gemm_desc));
}
}
}
break;
case miopenBFloat16: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
for(int i = 0; i < gemm_desc.batch_count; ++i)
{
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const rocblas_bfloat16*>(A) + a_offset + i * gemm_desc.strideA,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.lda,
static_cast<const rocblas_bfloat16*>(B) + b_offset + i * gemm_desc.strideB,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldb,
&beta,
static_cast<const rocblas_bfloat16*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
static_cast<rocblas_half*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_bf16_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
}
break;
case miopenFloat: {
float alpha = gemm_desc.alpha;
float beta = gemm_desc.beta;
for(int i = 0; i < gemm_desc.batch_count; ++i)
{
rb_status = miopen_rocblas_gemm_ex(
handle,
gemm_desc,
gemm_desc.transA ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.transB ? rocblas_operation_transpose : rocblas_operation_none,
gemm_desc.m,
gemm_desc.n,
gemm_desc.k,
&alpha,
static_cast<const float*>(A) + a_offset + i * gemm_desc.strideA,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.lda,
static_cast<const float*>(B) + b_offset + i * gemm_desc.strideB,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldb,
&beta,
static_cast<const float*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
static_cast<float*>(C) + c_offset + i * gemm_desc.strideC,
rocblas_datatype::rocblas_datatype_f32_r,
gemm_desc.ldc,
rocBlasComputeType(gemm_desc), // rocblas_datatype::rocblas_datatype_f32_r,
rocblas_gemm_algo::rocblas_gemm_algo_standard,
0,
0);
}
}
break;
case miopenFloat8_fnuz:
case miopenBFloat8_fnuz: {
const auto is_gfx942 = handle.GetDeviceName() == "gfx942";
if(is_gfx942)
{
rb_status = miopen_rocblas_gemm_strided_batched_ex3<char>(
handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset);
}
else
{
MIOPEN_THROW(miopenStatusInternalError,
"8-bit floating types are only supported on gfx942");
}
break;
}
case miopenDouble: {
MIOPEN_THROW(miopenStatusInternalError,
"miopenDouble data type not supported by rocBLAS.");
}
break;
case miopenInt64: {
MIOPEN_THROW(miopenStatusInternalError, "miopenInt64 is not currently supported.");
}
break;
}
if(handle.IsProfilingEnabled())
ProfilingRecordStop(handle, start, stop);
if(rb_status != rocblas_status::rocblas_status_success)
MIOPEN_THROW(miopenStatusInternalError, "rocBlas error encountered");
if(gemm_desc.deterministic)
SetRocblasAtomics(handle, cur_mode);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
case GemmBackend_t::hipblaslt: {
#if MIOPEN_USE_HIPBLASLT
// todo bharriso - find out if we need to support iterative variant, or if using regular
// batching is alright.
call_miopen_hipblasLt_gemm(handle, gemm_desc, A, a_offset, B, b_offset, C, c_offset, false);
return miopenStatusSuccess;
#else
return miopenStatusNotImplemented;
#endif
}
}
return miopenStatusUnknownError;
}
// y = w * Im2Col(x)
GemmDescriptor CreateGemmDescriptorConvFwd(const TensorDescriptor& wDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& yDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == xDesc.GetType());
if(wDesc.GetType() != miopenInt8)
assert(wDesc.GetType() == yDesc.GetType());
#endif
int in_c = xDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto wei_spatial =
wDesc.GetLengths() | std::views::drop(2) | std::views::take(wDesc.GetLengths().size() - 2);
auto out_spatial =
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = (wDesc.GetType() == miopenInt8);
int m = wei_k;
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k =
in_c * std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int lda = k;
int ldb = wDesc.GetType() == miopenInt8 ? k : n;
int ldc = n;
int batch_count = 1;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(0);
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// dx = Col2Im(transpose(w) * dy)
GemmDescriptor CreateGemmDescriptorConvBwdData(const TensorDescriptor& wDesc,
const TensorDescriptor& dyDesc,
const TensorDescriptor& dxDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
#endif
int in_c = dxDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto wei_spatial =
wDesc.GetLengths() | std::views::drop(2) | std::views::take(wDesc.GetLengths().size() - 2);
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = true;
bool transB = false;
int m =
in_c * std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = wei_k;
int lda = m;
int ldb = n;
int ldc = n;
int batch_count = 1;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(0);
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
dxDesc.GetType(),
false};
}
// dw = dy * transpose(Im2Col(x))
GemmDescriptor CreateGemmDescriptorConvBwdWeight(const TensorDescriptor& dyDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& dwDesc)
{
#ifndef NDEBUG
assert(dwDesc.GetType() == xDesc.GetType() && dwDesc.GetType() == dyDesc.GetType());
#endif
std::size_t in_c = xDesc.GetLengths()[1];
std::size_t wei_k = dwDesc.GetLengths()[0];
auto wei_spatial = dwDesc.GetLengths() | std::views::drop(2) |
std::views::take(dwDesc.GetLengths().size() - 2);
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = true;
int m = wei_k;
int n = static_cast<int>(in_c) *
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int k = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int lda = k;
int ldb = k;
int ldc = n;
int batch_count = 1;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(0);
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 1.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// y = CNHW2NCHW(w * NCHW2CNHW(x))
GemmDescriptor CreateGemmDescriptorConvCNHWFwd(const TensorDescriptor& wDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& yDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == xDesc.GetType());
if(wDesc.GetType() != miopenInt8)
assert(wDesc.GetType() == yDesc.GetType());
#endif
int in_n = xDesc.GetLengths()[0];
int in_c = xDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto out_spatial =
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = (wDesc.GetType() == miopenInt8);
int m = wei_k;
int n =
in_n * std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = in_c;
int lda = k;
int ldb = wDesc.GetType() == miopenInt8 ? k : n;
int ldc = n;
int batch_count = 1;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(0);
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// dx = CNHW2NCHW(transpose(w) * NCHW2CNHW(dy))
GemmDescriptor CreateGemmDescriptorConvCNHWBwdData(const TensorDescriptor& wDesc,
const TensorDescriptor& dyDesc,
const TensorDescriptor& dxDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
#endif
int in_n = dxDesc.GetLengths()[0];
int in_c = dxDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = true;
bool transB = false;
int m = in_c;
int n =
in_n * std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = wei_k;
int lda = m;
int ldb = n;
int ldc = n;
int batch_count = 1;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(0);
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
dxDesc.GetType(),
false};
}
// y[i] = w * x[i], i is batch id
GemmDescriptor CreateGemmStridedBatchedDescriptorConv1x1Fwd(const TensorDescriptor& wDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& yDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == xDesc.GetType());
if(wDesc.GetType() != miopenInt8)
assert(wDesc.GetType() == yDesc.GetType());
#else
(void)yDesc;
#endif
int in_n = xDesc.GetLengths()[0];
int in_c = xDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto in_spatial =
xDesc.GetLengths() | std::views::drop(2) | std::views::take(xDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = (wDesc.GetType() == miopenInt8);
int m = wei_k;
int n = std::accumulate(in_spatial.begin(), in_spatial.end(), 1, std::multiplies<int>());
int k = in_c;
int lda = k;
int ldb = wDesc.GetType() == miopenInt8 ? k : n;
int ldc = n;
int batch_count = in_n;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// dx[i] = transpose(w) * dy[i], i is batch id
GemmDescriptor CreateGemmStridedBatchedDescriptorConv1x1BwdData(const TensorDescriptor& wDesc,
const TensorDescriptor& dyDesc,
const TensorDescriptor& dxDesc)
{
#ifndef NDEBUG
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
#else
(void)dyDesc;
#endif
int in_n = dxDesc.GetLengths()[0];
int in_c = dxDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto in_spatial = dxDesc.GetLengths() | std::views::drop(2) |
std::views::take(dxDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = true;
bool transB = false;
int m = in_c;
int n = std::accumulate(in_spatial.begin(), in_spatial.end(), 1, std::multiplies<int>());
int k = wei_k;
int lda = m;
int ldb = n;
int ldc = n;
int batch_count = in_n;
auto strideA = static_cast<long long>(0);
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
dxDesc.GetType(),
false};
}
// dw = sum_over_batch(dy[i] * transpose(x[i])), i is batch id
GemmDescriptor CreateGemmStridedBatchedDescriptorConv1x1BwdWeight(const TensorDescriptor& dyDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& dwDesc)
{
#ifndef NDEBUG
assert(dwDesc.GetType() == xDesc.GetType() && dwDesc.GetType() == dyDesc.GetType());
#else
(void)dyDesc;
#endif
int in_n = xDesc.GetLengths()[0];
int in_c = xDesc.GetLengths()[1];
int wei_k = dwDesc.GetLengths()[0];
auto in_spatial =
xDesc.GetLengths() | std::views::drop(2) | std::views::take(xDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = true;
int m = wei_k;
int n = in_c;
int k = std::accumulate(in_spatial.begin(), in_spatial.end(), 1, std::multiplies<int>());
int lda = k;
int ldb = k;
int ldc = n;
int batch_count = in_n;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(0);
float alpha = 1.;
float beta = 1.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// y = w * Im2Col(x)
GemmDescriptor CreateGemmDescriptorGroupConvFwd(const TensorDescriptor& wDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& yDesc,
int groupCount)
{
#ifndef NDEBUG
assert(wDesc.GetType() == xDesc.GetType() && wDesc.GetType() == yDesc.GetType());
#endif
int in_c = xDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto wei_spatial =
wDesc.GetLengths() | std::views::drop(2) | std::views::take(wDesc.GetLengths().size() - 2);
auto out_spatial =
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = false;
int m = wei_k / groupCount;
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = (in_c / groupCount) *
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int lda = k;
int ldb = n;
int ldc = n;
int batch_count = groupCount;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// dx = Col2Im(transpose(w) * dy)
GemmDescriptor CreateGemmDescriptorGroupConvBwdData(const TensorDescriptor& wDesc,
const TensorDescriptor& dyDesc,
const TensorDescriptor& dxDesc,
int groupCount)
{
#ifndef NDEBUG
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
#endif
int in_c = dxDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto wei_spatial =
wDesc.GetLengths() | std::views::drop(2) | std::views::take(wDesc.GetLengths().size() - 2);
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = true;
bool transB = false;
int m = (in_c / groupCount) *
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int n = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = wei_k / groupCount;
int lda = m;
int ldb = n;
int ldc = n;
int batch_count = groupCount;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
dxDesc.GetType(),
false};
}
// dw = dy * transpose(Im2Col(x))
GemmDescriptor CreateGemmDescriptorGroupConvBwdWeight(const TensorDescriptor& dyDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& dwDesc,
int groupCount)
{
#ifndef NDEBUG
assert(dwDesc.GetType() == xDesc.GetType() && dwDesc.GetType() == dyDesc.GetType());
#endif
int in_c = xDesc.GetLengths()[1];
int wei_k = dwDesc.GetLengths()[0];
auto wei_spatial = dwDesc.GetLengths() | std::views::drop(2) |
std::views::take(dwDesc.GetLengths().size() - 2);
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = true;
int m = wei_k / groupCount;
int n = (in_c / groupCount) *
std::accumulate(wei_spatial.begin(), wei_spatial.end(), 1, std::multiplies<int>());
int k = std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int lda = k;
int ldb = k;
int ldc = n;
int batch_count = groupCount;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 1.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// y = CNHW2NCHW(w * NCHW2CNHW(x))
GemmDescriptor CreateGemmDescriptorGroupConvCNHWFwd(const TensorDescriptor& wDesc,
const TensorDescriptor& xDesc,
const TensorDescriptor& yDesc,
int groupCount)
{
#ifndef NDEBUG
assert(wDesc.GetType() == xDesc.GetType() && wDesc.GetType() == yDesc.GetType());
#endif
int in_n = xDesc.GetLengths()[0];
int in_c = xDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto out_spatial =
yDesc.GetLengths() | std::views::drop(2) | std::views::take(yDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = false;
bool transB = false;
int m = wei_k / groupCount;
int n =
in_n * std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = in_c / groupCount;
int lda = k;
int ldb = n;
int ldc = n;
int batch_count = groupCount;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
xDesc.GetType(),
false};
}
// dx = CNHW2NCHW(transpose(w) * NCHW2CNHW(dy))
GemmDescriptor CreateGemmDescriptorGroupConvCNHWBwdData(const TensorDescriptor& wDesc,
const TensorDescriptor& dyDesc,
const TensorDescriptor& dxDesc,
int groupCount)
{
#ifndef NDEBUG
assert(wDesc.GetType() == dxDesc.GetType() && wDesc.GetType() == dyDesc.GetType());
#endif
int in_n = dxDesc.GetLengths()[0];
int in_c = dxDesc.GetLengths()[1];
int wei_k = wDesc.GetLengths()[0];
auto out_spatial = dyDesc.GetLengths() | std::views::drop(2) |
std::views::take(dyDesc.GetLengths().size() - 2);
bool isColMajor = false;
bool transA = true;
bool transB = false;
int m = in_c / groupCount;
int n =
in_n * std::accumulate(out_spatial.begin(), out_spatial.end(), 1, std::multiplies<int>());
int k = wei_k / groupCount;
int lda = m;
int ldb = n;
int ldc = n;
int batch_count = groupCount;
auto strideA = static_cast<long long>(m) * k;
auto strideB = static_cast<long long>(k) * n;
auto strideC = static_cast<long long>(m) * n;
float alpha = 1.;
float beta = 0.;
return GemmDescriptor{isColMajor,
transA,
transB,
m,
n,
k,
lda,
ldb,
ldc,
batch_count,
strideA,
strideB,
strideC,
alpha,
beta,
dxDesc.GetType(),
false};
}
} // namespace miopen
Reference in new issue View Git Blame Copy Permalink