derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
515 lines
17 KiB
C++
515 lines
17 KiB
C++
/*******************************************************************************
|
|
*
|
|
* MIT License
|
|
*
|
|
* Copyright (C) 2021-2025 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
* of this software and associated documentation files (the "Software"), to deal
|
|
* in the Software without restriction, including without limitation the rights
|
|
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
* copies of the Software, and to permit persons to whom the Software is
|
|
* furnished to do so, subject to the following conditions:
|
|
*
|
|
* The above copyright notice and this permission notice shall be included in
|
|
* all copies or substantial portions of the Software.
|
|
*
|
|
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
* SOFTWARE.
|
|
*
|
|
*******************************************************************************/
|
|
#ifndef ROCWMMA_ROCBLAS_REFERENCE_HPP
|
|
#define ROCWMMA_ROCBLAS_REFERENCE_HPP
|
|
|
|
#define ROCM_USE_FLOAT16
|
|
|
|
#define CHECK_ROCBLAS_ERROR(expression) \
|
|
if(auto status = (expression); status != rocblas_status_success) \
|
|
{ \
|
|
fprintf(stderr, \
|
|
"rocBLAS error: '%s'(%d) at %s:%d\n", \
|
|
rocblas_status_to_string(status), \
|
|
status, \
|
|
__FILE__, \
|
|
__LINE__); \
|
|
exit(EXIT_FAILURE); \
|
|
}
|
|
|
|
// BETA_FEATURES_API needs to be defined to access the beta features of rocBLAS which includes float8/bfloat8 support.
|
|
#define ROCBLAS_BETA_FEATURES_API
|
|
#include <rocblas/rocblas.h>
|
|
|
|
#if(ROCBLAS_VERSION_MAJOR >= 3) || ((ROCBLAS_VERSION_MAJOR >= 2) && (ROCBLAS_VERSION_MINOR > 45))
|
|
#define ROCBLAS_DATA_TYPE_INVALID
|
|
#endif
|
|
|
|
#if(((ROCBLAS_VERSION_MAJOR == 3) && (ROCBLAS_VERSION_MINOR >= 1)) || (ROCBLAS_VERSION_MAJOR == 4))
|
|
#define ROCBLAS_DATA_TYPE_FLOAT8
|
|
#endif
|
|
|
|
#include "common.hpp"
|
|
#include <rocwmma/internal/types.hpp>
|
|
|
|
namespace rocwmma
|
|
{
|
|
|
|
constexpr const char* rocblas_datatype2string(rocblas_datatype type)
|
|
{
|
|
switch(type)
|
|
{
|
|
case rocblas_datatype_f16_r:
|
|
return "f16_r";
|
|
case rocblas_datatype_f32_r:
|
|
return "f32_r";
|
|
case rocblas_datatype_f64_r:
|
|
return "f64_r";
|
|
case rocblas_datatype_f16_c:
|
|
return "f16_c";
|
|
case rocblas_datatype_f32_c:
|
|
return "f32_c";
|
|
case rocblas_datatype_f64_c:
|
|
return "f64_c";
|
|
case rocblas_datatype_i8_r:
|
|
return "i8_r";
|
|
case rocblas_datatype_u8_r:
|
|
return "u8_r";
|
|
case rocblas_datatype_i32_r:
|
|
return "i32_r";
|
|
case rocblas_datatype_u32_r:
|
|
return "u32_r";
|
|
case rocblas_datatype_i8_c:
|
|
return "i8_c";
|
|
case rocblas_datatype_u8_c:
|
|
return "u8_c";
|
|
case rocblas_datatype_i32_c:
|
|
return "i32_c";
|
|
case rocblas_datatype_u32_c:
|
|
return "u32_c";
|
|
case rocblas_datatype_bf16_r:
|
|
return "bf16_r";
|
|
case rocblas_datatype_bf16_c:
|
|
return "bf16_c";
|
|
#if defined(ROCBLAS_DATA_TYPE_FLOAT8)
|
|
case rocblas_datatype_f8_r:
|
|
return "f8_r";
|
|
case rocblas_datatype_bf8_r:
|
|
return "bf8_r";
|
|
#endif
|
|
#if defined(ROCBLAS_DATA_TYPE_INVALID)
|
|
case rocblas_datatype_invalid:;
|
|
#endif // ROCBLAS_DATA_TYPE_INVALID
|
|
}
|
|
return "invalid";
|
|
}
|
|
|
|
template <typename DataT>
|
|
struct rocblas_types
|
|
#if defined(ROCBLAS_DATA_TYPE_INVALID)
|
|
{
|
|
using DataType = DataT;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_invalid;
|
|
}
|
|
}
|
|
#endif //ROCBLAS_DATA_TYPE_INVALID
|
|
;
|
|
|
|
template <>
|
|
struct rocblas_types<int8_t>
|
|
{
|
|
using DataType = int8_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_i8_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<uint8_t>
|
|
{
|
|
using DataType = uint8_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_u8_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<int32_t>
|
|
{
|
|
using DataType = rocblas_int;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_i32_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<uint32_t>
|
|
{
|
|
using DataType = uint32_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_u32_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<float16_t>
|
|
{
|
|
using DataType = rocblas_half;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f16_r;
|
|
}
|
|
};
|
|
|
|
#if !ROCWMMA_TESTS_NO_HALF
|
|
template <>
|
|
struct rocblas_types<hfloat16_t> : public rocblas_types<float16_t>
|
|
{
|
|
};
|
|
#endif // !ROCWMMA_TESTS_NO_HALF
|
|
|
|
template <>
|
|
struct rocblas_types<bfloat16_t>
|
|
{
|
|
using DataType = rocblas_bfloat16;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_bf16_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<float32_t>
|
|
{
|
|
using DataType = float32_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f32_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<xfloat32_t>
|
|
{
|
|
using DataType = xfloat32_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f32_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<float64_t>
|
|
{
|
|
using DataType = float64_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f64_r;
|
|
}
|
|
};
|
|
|
|
#if defined(ROCBLAS_DATA_TYPE_FLOAT8)
|
|
template <>
|
|
struct rocblas_types<bfloat8_t>
|
|
{
|
|
using DataType = bfloat8_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_bf8_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<float8_t>
|
|
{
|
|
using DataType = float8_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f8_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<bfloat8_fnuz_t>
|
|
{
|
|
using DataType = bfloat8_fnuz_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_bf8_r;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_types<float8_fnuz_t>
|
|
{
|
|
using DataType = float8_fnuz_t;
|
|
constexpr static inline rocblas_datatype type()
|
|
{
|
|
return rocblas_datatype_f8_r;
|
|
}
|
|
};
|
|
#endif
|
|
|
|
template <typename DataLayoutT>
|
|
struct rocblas_layout;
|
|
|
|
template <>
|
|
struct rocblas_layout<row_major>
|
|
{
|
|
using Layout = row_major;
|
|
constexpr static inline rocblas_operation operation()
|
|
{
|
|
return rocblas_operation_transpose;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct rocblas_layout<col_major>
|
|
{
|
|
using Layout = col_major;
|
|
constexpr static inline rocblas_operation operation()
|
|
{
|
|
return rocblas_operation_none;
|
|
}
|
|
};
|
|
|
|
//! @brief Dispatcher for rocBLAS.
|
|
//! Need to call different interface for f8 validation
|
|
auto dispatch_rocBLAS(rocblas_handle handle,
|
|
rocblas_operation opA,
|
|
rocblas_operation opB,
|
|
rocblas_int m,
|
|
rocblas_int n,
|
|
rocblas_int k,
|
|
const void* alpha,
|
|
const void* a,
|
|
rocblas_datatype a_type,
|
|
rocblas_int lda,
|
|
const void* b,
|
|
rocblas_datatype b_type,
|
|
rocblas_int ldb,
|
|
const void* beta,
|
|
const void* c,
|
|
rocblas_datatype c_type,
|
|
rocblas_int ldc,
|
|
void* d,
|
|
rocblas_datatype d_type,
|
|
rocblas_int ldd,
|
|
rocblas_datatype compute_type,
|
|
rocblas_gemm_algo algo,
|
|
int32_t solution_index,
|
|
uint32_t flags)
|
|
{
|
|
#if defined(ROCBLAS_DATA_TYPE_FLOAT8)
|
|
if(a_type == rocblas_datatype_f8_r || b_type == rocblas_datatype_f8_r
|
|
|| a_type == rocblas_datatype_bf8_r || b_type == rocblas_datatype_bf8_r)
|
|
{
|
|
if(compute_type != rocblas_datatype_f32_r)
|
|
{
|
|
std::cerr << "float8_t and bfloat8_t input types must have f32 compute type with "
|
|
"this version of rocBLAS"
|
|
<< std::endl;
|
|
return rocblas_status_not_implemented;
|
|
}
|
|
|
|
// rocblas_gemm_ex3 needs a rocblas_computetype
|
|
rocblas_computetype compute_type_f32 = rocblas_compute_type_f32;
|
|
if(a_type == rocblas_datatype_f8_r && b_type == rocblas_datatype_f8_r)
|
|
{
|
|
compute_type_f32 = rocblas_compute_type_f8_f8_f32;
|
|
}
|
|
else if(a_type == rocblas_datatype_f8_r && b_type == rocblas_datatype_bf8_r)
|
|
{
|
|
compute_type_f32 = rocblas_compute_type_f8_bf8_f32;
|
|
}
|
|
else if(a_type == rocblas_datatype_bf8_r && b_type == rocblas_datatype_f8_r)
|
|
{
|
|
compute_type_f32 = rocblas_compute_type_bf8_f8_f32;
|
|
}
|
|
else if(a_type == rocblas_datatype_bf8_r && b_type == rocblas_datatype_bf8_r)
|
|
{
|
|
compute_type_f32 = rocblas_compute_type_bf8_bf8_f32;
|
|
}
|
|
|
|
return rocblas_gemm_ex3(handle,
|
|
opA,
|
|
opB,
|
|
m,
|
|
n,
|
|
k,
|
|
alpha,
|
|
a,
|
|
a_type,
|
|
lda,
|
|
b,
|
|
b_type,
|
|
ldb,
|
|
beta,
|
|
c,
|
|
c_type,
|
|
ldc,
|
|
d,
|
|
d_type,
|
|
ldd,
|
|
compute_type_f32,
|
|
algo,
|
|
solution_index,
|
|
flags);
|
|
}
|
|
#endif
|
|
|
|
return rocblas_gemm_ex(handle,
|
|
opA,
|
|
opB,
|
|
m,
|
|
n,
|
|
k,
|
|
alpha,
|
|
a,
|
|
a_type,
|
|
lda,
|
|
b,
|
|
b_type,
|
|
ldb,
|
|
beta,
|
|
c,
|
|
c_type,
|
|
ldc,
|
|
d,
|
|
d_type,
|
|
ldd,
|
|
compute_type,
|
|
algo,
|
|
solution_index,
|
|
flags);
|
|
}
|
|
/*
|
|
* Rocblas notes:
|
|
* Layouts C and D are always assumed as col_major
|
|
* Non-transpose (N) is col-major
|
|
* Transpose (T) is row-major
|
|
*/
|
|
template <typename InputT,
|
|
typename OutputT,
|
|
typename ComputeT,
|
|
typename LayoutA,
|
|
typename LayoutB>
|
|
void gemm_rocBLAS(uint32_t m,
|
|
uint32_t n,
|
|
uint32_t k,
|
|
InputT const* ha,
|
|
InputT const* hb,
|
|
OutputT const* hc,
|
|
OutputT* hd,
|
|
ComputeT alpha,
|
|
ComputeT beta)
|
|
{
|
|
rocblas_datatype a_type = rocblas_types<InputT>::type();
|
|
rocblas_datatype b_type = rocblas_types<InputT>::type();
|
|
rocblas_datatype c_type = rocblas_types<OutputT>::type();
|
|
rocblas_datatype d_type = rocblas_types<OutputT>::type();
|
|
rocblas_datatype compute_type = rocblas_types<ComputeT>::type();
|
|
|
|
using a_t = typename rocblas_types<InputT>::DataType;
|
|
using b_t = typename rocblas_types<InputT>::DataType;
|
|
using c_t = typename rocblas_types<OutputT>::DataType;
|
|
using d_t = typename rocblas_types<OutputT>::DataType;
|
|
|
|
size_t size_a = m * k;
|
|
size_t size_b = k * n;
|
|
size_t size_c = m * n;
|
|
size_t size_d = m * n;
|
|
|
|
rocblas_operation opA = rocblas_layout<LayoutA>::operation();
|
|
rocblas_operation opB = rocblas_layout<LayoutB>::operation();
|
|
|
|
rocblas_int lda = (opA == rocblas_operation_none ? m : k);
|
|
rocblas_int ldb = (opB == rocblas_operation_none ? k : n);
|
|
rocblas_int ldc = m;
|
|
rocblas_int ldd = m;
|
|
|
|
// allocate memory on device
|
|
a_t* da;
|
|
b_t* db;
|
|
c_t* dc;
|
|
d_t* dd;
|
|
|
|
CHECK_HIP_ERROR(hipMalloc(&da, size_a * sizeof(a_t)));
|
|
CHECK_HIP_ERROR(hipMalloc(&db, size_b * sizeof(b_t)));
|
|
CHECK_HIP_ERROR(hipMalloc(&dc, size_c * sizeof(c_t)));
|
|
CHECK_HIP_ERROR(hipMalloc(&dd, size_d * sizeof(d_t)));
|
|
|
|
// copy matrices from host to device
|
|
CHECK_HIP_ERROR(hipMemcpy(da, &ha[0], sizeof(a_t) * size_a, hipMemcpyHostToDevice));
|
|
CHECK_HIP_ERROR(hipMemcpy(db, &hb[0], sizeof(b_t) * size_b, hipMemcpyHostToDevice));
|
|
CHECK_HIP_ERROR(hipMemcpy(dc, &hc[0], sizeof(c_t) * size_c, hipMemcpyHostToDevice));
|
|
CHECK_HIP_ERROR(hipMemcpy(dd, &hd[0], sizeof(d_t) * size_d, hipMemcpyHostToDevice));
|
|
|
|
rocblas_handle handle;
|
|
CHECK_ROCBLAS_ERROR(rocblas_create_handle(&handle));
|
|
|
|
auto algo = rocblas_gemm_algo_standard;
|
|
int32_t solution_index = 0;
|
|
uint32_t flags = 0;
|
|
|
|
if((std::is_same<InputT, float8_t>::value)
|
|
|| (std::is_same<InputT, bfloat8_t>::value
|
|
|| (std::is_same<InputT, float8_fnuz_t>::value)
|
|
|| (std::is_same<InputT, bfloat8_fnuz_t>::value)))
|
|
{
|
|
#if !defined(ROCBLAS_DATA_TYPE_FLOAT8)
|
|
std::cerr
|
|
<< "float8_t and bfloat8_t types are not supported with this version of rocBLAS"
|
|
<< std::endl;
|
|
return;
|
|
#endif
|
|
}
|
|
|
|
CHECK_ROCBLAS_ERROR(dispatch_rocBLAS(handle,
|
|
opA,
|
|
opB,
|
|
m,
|
|
n,
|
|
k,
|
|
&alpha,
|
|
da,
|
|
a_type,
|
|
lda,
|
|
db,
|
|
b_type,
|
|
ldb,
|
|
&beta,
|
|
dc,
|
|
c_type,
|
|
ldc,
|
|
dd,
|
|
d_type,
|
|
ldd,
|
|
compute_type,
|
|
algo,
|
|
solution_index,
|
|
flags));
|
|
|
|
// copy output from device to CPU
|
|
CHECK_HIP_ERROR(hipMemcpy(&hd[0], dd, sizeof(d_t) * size_d, hipMemcpyDeviceToHost));
|
|
|
|
CHECK_HIP_ERROR(hipFree(da));
|
|
CHECK_HIP_ERROR(hipFree(db));
|
|
CHECK_HIP_ERROR(hipFree(dc));
|
|
CHECK_HIP_ERROR(hipFree(dd));
|
|
CHECK_ROCBLAS_ERROR(rocblas_destroy_handle(handle));
|
|
}
|
|
|
|
} // namespace rocwmma
|
|
|
|
#endif // ROCWMMA_ROCBLAS_REFERENCE_HPP
|