saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] Refactor to hide ck layout (Row/Col) from ORT interface (#18777)

Browse source 
Previously, we use `ck::tensor_layout::gemm::RowMajor` or `ColumnMajor`
to tag the template for correct dispatch. This is cumbersome in the case
of CK is disabled.

Switch to use the ORT BlasOp to tag the template and use
`CKBlasOpAdaptor` to adapt between ORT BlasOp enum and ck's Col/Row.
Just like what we have done for ORT datatype and ck datatype with
`CKDataTypeAdaptor`.
This commit is contained in:
cloudhan 2023-12-13 11:37:26 +08:00 committed by GitHub
parent 0ca84549ab
commit 3940ef20be
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 252 additions and 245 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh
View file

@ -31,6 +31,7 @@ namespace internal {
#ifdef USE_COMPOSABLE_KERNEL
using onnxruntime::rocm::CKDataTypeAdaptor;
using onnxruntime::rocm::CKBlasOpAdaptor;
using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;
@ -39,9 +40,11 @@ using Nop = ck::tensor_operation::element_wise::PassThrough;
using AddFastGelu = ck::tensor_operation::element_wise::AddFastGelu;
using FastGelu = ck::tensor_operation::element_wise::FastGelu;
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKGemmAddFastGeluTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceGemmAddFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD<
ALayout, BLayout, ck::Tuple<Row>, Row,
CKDataType, CKDataType, ck::Tuple<CKDataType>, CKDataType,
@ -76,9 +79,11 @@ auto GetCKGemmAddFastGeluTypeStringAndOps() {
return ret;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKGemmFastGeluTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceGemmFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD<
ALayout, BLayout, ck::Tuple<>, Row,
CKDataType, CKDataType, ck::Tuple<>, CKDataType,

8
onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_impl.cu
View file

@ -49,16 +49,16 @@ inline GEMMFASTGELU(T, ScalarT) {
if (tuning_ctx->IsTunableOpEnabled()) {
if (opa == BlasOp::N && opb == BlasOp::N) {
static internal::GemmFastGeluTunableOp<T, internal::Row, internal::Row> gemm_fast_gelu{};
static internal::GemmFastGeluTunableOp<T, BlasOp::N, BlasOp::N> gemm_fast_gelu{};
return gemm_fast_gelu(&params);
} else if (opa == BlasOp::T && opb == BlasOp::N) {
static internal::GemmFastGeluTunableOp<T, internal::Col, internal::Row> gemm_fast_gelu{};
static internal::GemmFastGeluTunableOp<T, BlasOp::T, BlasOp::N> gemm_fast_gelu{};
return gemm_fast_gelu(&params);
} else if (opa == BlasOp::N && opb == BlasOp::T) {
static internal::GemmFastGeluTunableOp<T, internal::Row, internal::Col> gemm_fast_gelu{};
static internal::GemmFastGeluTunableOp<T, BlasOp::N, BlasOp::T> gemm_fast_gelu{};
return gemm_fast_gelu(&params);
} else /*if (opa == BlasOp::T && opb == BlasOp::T)*/ {
static internal::GemmFastGeluTunableOp<T, internal::Col, internal::Col> gemm_fast_gelu{};
static internal::GemmFastGeluTunableOp<T, BlasOp::T, BlasOp::T> gemm_fast_gelu{};
return gemm_fast_gelu(&params);
}
}

8
onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh
View file

@ -51,24 +51,24 @@ Status GemmFastGeluUnfused(const GemmFastGeluParams<T>* params) {
params->c);
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmFastGeluTunableOp : public TunableOp<GemmFastGeluParams<T>> {
public:
GemmFastGeluTunableOp() {
this->RegisterOp(GemmFastGeluUnfused<T>);
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKGemmAddFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKGemmAddFastGeluTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKGemmFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKGemmFastGeluTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
#endif
#ifdef USE_HIPBLASLT
for (auto&& [_, op] : GetHipBlasLtGemmFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetHipBlasLtGemmFastGeluTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}

24
onnxruntime/core/providers/rocm/tunable/gemm.cu
View file

@ -53,16 +53,16 @@ inline GEMM(T, ScalarT) {
if (tuning_ctx->IsTunableOpEnabled()) {
if (opa == BlasOp::N && opb == BlasOp::N) {
static internal::GemmTunableOp<T, internal::Row, internal::Row> gemm{};
static internal::GemmTunableOp<T, BlasOp::N, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::T && opb == BlasOp::N) {
static internal::GemmTunableOp<T, internal::Col, internal::Row> gemm{};
static internal::GemmTunableOp<T, BlasOp::T, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::N && opb == BlasOp::T) {
static internal::GemmTunableOp<T, internal::Row, internal::Col> gemm{};
static internal::GemmTunableOp<T, BlasOp::N, BlasOp::T> gemm{};
return gemm(&params);
} else /*if (opa == BlasOp::T && opb == BlasOp::T)*/ {
static internal::GemmTunableOp<T, internal::Col, internal::Col> gemm{};
static internal::GemmTunableOp<T, BlasOp::T, BlasOp::T> gemm{};
return gemm(&params);
}
}
@ -94,16 +94,16 @@ inline BATCHED_GEMM(T, ScalarT) {
if (tuning_ctx->IsTunableOpEnabled()) {
if (opa == BlasOp::N && opb == BlasOp::N) {
static internal::BatchedGemmTunableOp<T, internal::Row, internal::Row> gemm{};
static internal::BatchedGemmTunableOp<T, BlasOp::N, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::T && opb == BlasOp::N) {
static internal::BatchedGemmTunableOp<T, internal::Col, internal::Row> gemm{};
static internal::BatchedGemmTunableOp<T, BlasOp::T, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::N && opb == BlasOp::T) {
static internal::BatchedGemmTunableOp<T, internal::Row, internal::Col> gemm{};
static internal::BatchedGemmTunableOp<T, BlasOp::N, BlasOp::T> gemm{};
return gemm(&params);
} else /*if (opa == BlasOp::T && opb == BlasOp::T)*/ {
static internal::BatchedGemmTunableOp<T, internal::Col, internal::Col> gemm{};
static internal::BatchedGemmTunableOp<T, BlasOp::T, BlasOp::T> gemm{};
return gemm(&params);
}
}
@ -138,16 +138,16 @@ inline STRIDED_BATCHED_GEMM(T, ScalarT) {
if (tuning_ctx->IsTunableOpEnabled()) {
if (opa == BlasOp::N && opb == BlasOp::N) {
static internal::StridedBatchedGemmTunableOp<T, internal::Row, internal::Row> gemm{};
static internal::StridedBatchedGemmTunableOp<T, BlasOp::N, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::T && opb == BlasOp::N) {
static internal::StridedBatchedGemmTunableOp<T, internal::Col, internal::Row> gemm{};
static internal::StridedBatchedGemmTunableOp<T, BlasOp::T, BlasOp::N> gemm{};
return gemm(&params);
} else if (opa == BlasOp::N && opb == BlasOp::T) {
static internal::StridedBatchedGemmTunableOp<T, internal::Row, internal::Col> gemm{};
static internal::StridedBatchedGemmTunableOp<T, BlasOp::N, BlasOp::T> gemm{};
return gemm(&params);
} else /*if (opa == BlasOp::T && opb == BlasOp::T)*/ {
static internal::StridedBatchedGemmTunableOp<T, internal::Col, internal::Col> gemm{};
static internal::StridedBatchedGemmTunableOp<T, BlasOp::T, BlasOp::T> gemm{};
return gemm(&params);
}
}

16
onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh
View file

@ -36,9 +36,11 @@ using Col = ck::tensor_layout::gemm::ColumnMajor;
using Nop = ck::tensor_operation::element_wise::PassThrough;
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKGemmTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceGemm = ck::tensor_operation::device::DeviceGemm<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,
@ -70,9 +72,11 @@ auto GetCKGemmTypeStringAndOps() {
return ret;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKStreamKGemmTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceGemm = ck::tensor_operation::device::DeviceGemmStreamK<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,
@ -104,9 +108,11 @@ auto GetCKStreamKGemmTypeStringAndOps() {
return ret;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKSplitKGemmTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceGemm = ck::tensor_operation::device::DeviceGemmSplitK<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,
@ -144,9 +150,11 @@ auto GetCKSplitKGemmTypeStringAndOps() {
return ret;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetCKStridedBatchedGemmTypeStringAndOps() {
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using ALayout = typename CKBlasOpAdaptor<OpA>::type;
using BLayout = typename CKBlasOpAdaptor<OpB>::type;
using DeviceStridedBatchedGemm = ck::tensor_operation::device::DeviceBatchedGemm<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,

24
onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h
View file

@ -59,9 +59,9 @@ constexpr hipblasltDatatype_t HipBlasDataTypeFor<double>() {
return HIPBLASLT_R_64F;
}
template <typename Layout>
constexpr hipblasOperation_t MapCKLayoutToHipBlasLt() {
if constexpr (std::is_same_v<Layout, Row>) {
template <BlasOp Op>
constexpr hipblasOperation_t MapBlasOpToHipBlasLt() {
if constexpr (Op == BlasOp::NonTrans) {
return HIPBLAS_OP_N;
}
return HIPBLAS_OP_T;
@ -101,13 +101,13 @@ std::string TypeStringFor() {
return "UnknownType";
}
template <typename T, typename ALayout, typename BLayout, typename ParamsT>
template <typename T, BlasOp OpA, BlasOp OpB, typename ParamsT>
auto GetHipBlasLtTypeStringAndOps(ActivationType activation_type = ActivationType::NONE) {
hipblasLtHandle_t handle;
HIPBLASLT_CALL_THROW(hipblasLtCreate(&handle));
hipblasOperation_t trans_a = MapCKLayoutToHipBlasLt<BLayout>();
hipblasOperation_t trans_b = MapCKLayoutToHipBlasLt<ALayout>();
hipblasOperation_t trans_a = MapBlasOpToHipBlasLt<OpB>();
hipblasOperation_t trans_b = MapBlasOpToHipBlasLt<OpA>();
hipblasltDatatype_t in_out_datatype = HipBlasDataTypeFor<T>();
std::vector<hipblasLtMatmulHeuristicResult_t> heuristic_result;
@ -266,19 +266,19 @@ auto GetHipBlasLtTypeStringAndOps(ActivationType activation_type = ActivationTyp
return ret;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetHipBlasLtGemmTypeStringAndOps() {
return GetHipBlasLtTypeStringAndOps<T, ALayout, BLayout, GemmParams<T>>();
return GetHipBlasLtTypeStringAndOps<T, OpA, OpB, GemmParams<T>>();
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetHipBlasLtStridedBatchedGemmTypeStringAndOps() {
return GetHipBlasLtTypeStringAndOps<T, ALayout, BLayout, StridedBatchedGemmParams<T>>();
return GetHipBlasLtTypeStringAndOps<T, OpA, OpB, StridedBatchedGemmParams<T>>();
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
auto GetHipBlasLtGemmFastGeluTypeStringAndOps() {
return GetHipBlasLtTypeStringAndOps<T, ALayout, BLayout, GemmFastGeluParams<T>>(ActivationType::GELU);
return GetHipBlasLtTypeStringAndOps<T, OpA, OpB, GemmFastGeluParams<T>>(ActivationType::GELU);
}
#endif // USE_HIPBLASLT

18
onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh
View file

@ -33,14 +33,14 @@ bool IsZero(half v) {
return __half2float(v) == 0.0f;
}
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmTunableOp : public TunableOp<GemmParams<T>> {
public:
GemmTunableOp() {
this->RegisterOp(RocBlasGemmOp<T>);
#ifdef USE_HIPBLASLT
for (auto&& [_, op] : GetHipBlasLtGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetHipBlasLtGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
@ -54,16 +54,16 @@ class GemmTunableOp : public TunableOp<GemmParams<T>> {
#endif
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKStreamKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKStreamKGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKSplitKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKSplitKGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
@ -96,7 +96,7 @@ class GemmTunableOp : public TunableOp<GemmParams<T>> {
}
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class BatchedGemmTunableOp : public TunableOp<BatchedGemmParams<T>> {
public:
BatchedGemmTunableOp() {
@ -146,14 +146,14 @@ class BatchedGemmTunableOp : public TunableOp<BatchedGemmParams<T>> {
}
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class StridedBatchedGemmTunableOp : public TunableOp<StridedBatchedGemmParams<T>> {
public:
StridedBatchedGemmTunableOp() {
this->RegisterOp(RocBlasStridedBatchedGemmOp<T>);
#ifdef USE_HIPBLASLT
for (auto&& [_, op] : GetHipBlasLtStridedBatchedGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetHipBlasLtStridedBatchedGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
@ -167,7 +167,7 @@ class StridedBatchedGemmTunableOp : public TunableOp<StridedBatchedGemmParams<T>
#endif
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKStridedBatchedGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [_, op] : GetCKStridedBatchedGemmTypeStringAndOps<T, OpA, OpB>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}

84
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu
View file

@ -23,7 +23,7 @@ namespace py = pybind11;
namespace onnxruntime {
#ifdef USE_COMPOSABLE_KERNEL
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class CKGemm : public IKernelExplorer {
public:
CKGemm(BlasOp opa, BlasOp opb,
@ -34,9 +34,7 @@ class CKGemm : public IKernelExplorer {
double beta,
DeviceArray& c, int64_t ldc)
: params_{} {
auto supports_a = opa == BlasOp::N ? std::is_same_v<ALayout, Row> : std::is_same_v<ALayout, Col>;
auto supports_b = opb == BlasOp::N ? std::is_same_v<BLayout, Row> : std::is_same_v<BLayout, Col>;
ORT_ENFORCE(supports_a && supports_b);
ORT_ENFORCE(opa == OpA && opb == OpB);
params_.tuning_ctx = TuningContext();
params_.stream = Stream();
@ -56,15 +54,15 @@ class CKGemm : public IKernelExplorer {
params_.c = static_cast<T*>(c.ptr());
params_.ldc = ldc;
for (auto&& [type_string, op] : GetCKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
for (auto&& [type_string, op] : GetCKStreamKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKStreamKGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
for (auto&& [type_string, op] : GetCKSplitKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKSplitKGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -100,7 +98,7 @@ class CKGemm : public IKernelExplorer {
size_t selected_op_{};
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class CKStridedBatchedGemm : public IKernelExplorer {
public:
CKStridedBatchedGemm(
@ -113,9 +111,7 @@ class CKStridedBatchedGemm : public IKernelExplorer {
DeviceArray& c, int64_t ldc, int64_t stride_c,
int64_t batch)
: params_{} {
auto supports_a = opa == BlasOp::N ? std::is_same_v<ALayout, Row> : std::is_same_v<ALayout, Col>;
auto supports_b = opb == BlasOp::N ? std::is_same_v<BLayout, Row> : std::is_same_v<BLayout, Col>;
ORT_ENFORCE(supports_a && supports_b);
ORT_ENFORCE(opa == OpA && opb == OpB);
params_.tuning_ctx = TuningContext();
params_.stream = Stream();
@ -139,7 +135,7 @@ class CKStridedBatchedGemm : public IKernelExplorer {
params_.stride_c = stride_c;
params_.batch = batch;
for (auto&& [type_string, op] : GetCKStridedBatchedGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKStridedBatchedGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -175,44 +171,44 @@ class CKStridedBatchedGemm : public IKernelExplorer {
size_t selected_op_{};
};
#define REGISTER_OP_COMMON(type, dtype, alayout, blayout, layout_string) \
py::class_<type<dtype, alayout, blayout>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, alayout, blayout>::SetRepeats) \
.def("Profile", &type<dtype, alayout, blayout>::Profile) \
.def("Run", &type<dtype, alayout, blayout>::Run) \
.def("ListOps", &type<dtype, alayout, blayout>::ListOps) \
.def("SelectOp", &type<dtype, alayout, blayout>::SelectOp)
#define REGISTER_OP_COMMON(type, dtype, opa, opb, layout_string) \
py::class_<type<dtype, opa, opb>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, opa, opb>::SetRepeats) \
.def("Profile", &type<dtype, opa, opb>::Profile) \
.def("Run", &type<dtype, opa, opb>::Run) \
.def("ListOps", &type<dtype, opa, opb>::ListOps) \
.def("SelectOp", &type<dtype, opa, opb>::SelectOp)
#define REGISTER_CKGEMM(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(CKGemm, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
#define REGISTER_CKGEMM(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(CKGemm, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
DeviceArray&, int64_t>());
#define REGISTER_CKGEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_CKGEMM(dtype, Row, Row, "NN"); \
REGISTER_CKGEMM(dtype, Row, Col, "NT"); \
REGISTER_CKGEMM(dtype, Col, Row, "TN"); \
REGISTER_CKGEMM(dtype, Col, Col, "TT");
#define REGISTER_CKGEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_CKGEMM(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_CKGEMM(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_CKGEMM(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_CKGEMM(dtype, BlasOp::T, BlasOp::T, "TT");
#define REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(CKStridedBatchedGemm, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
#define REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(CKStridedBatchedGemm, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
int64_t>());
#define REGISTER_CKSTRIDEDBATCHEDGEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, Row, Row, "NN"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, Row, Col, "NT"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, Col, Row, "TN"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, Col, Col, "TT");
#define REGISTER_CKSTRIDEDBATCHEDGEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_CKSTRIDEDBATCHEDGEMM(dtype, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_CKGEMM_FOR_ALL_TRANSAB(float);

48
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu
View file

@ -23,7 +23,7 @@ namespace py = pybind11;
namespace onnxruntime {
#ifdef USE_COMPOSABLE_KERNEL
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class CKGemmFastGelu : public IKernelExplorer {
public:
CKGemmFastGelu(BlasOp opa, BlasOp opb,
@ -35,9 +35,7 @@ class CKGemmFastGelu : public IKernelExplorer {
double beta,
DeviceArray& c, int64_t ldc)
: params_{} {
auto supports_a = opa == BlasOp::N ? std::is_same_v<ALayout, Row> : std::is_same_v<ALayout, Col>;
auto supports_b = opb == BlasOp::N ? std::is_same_v<BLayout, Row> : std::is_same_v<BLayout, Col>;
ORT_ENFORCE(supports_a && supports_b);
ORT_ENFORCE(opa == OpA && opb == OpB);
params_.tuning_ctx = TuningContext();
params_.stream = Stream();
@ -58,11 +56,11 @@ class CKGemmFastGelu : public IKernelExplorer {
params_.c = static_cast<T*>(c.ptr());
params_.ldc = ldc;
for (auto&& [type_string, op] : GetCKGemmAddFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKGemmAddFastGeluTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
for (auto&& [type_string, op] : GetCKGemmFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetCKGemmFastGeluTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -97,26 +95,26 @@ class CKGemmFastGelu : public IKernelExplorer {
size_t selected_op_{};
};
#define REGISTER_OP(type, alayout, blayout, layout_string) \
py::class_<CKGemmFastGelu<type, alayout, blayout>>(m, "CKGemmFastGelu_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &CKGemmFastGelu<type, alayout, blayout>::SetRepeats) \
.def("Profile", &CKGemmFastGelu<type, alayout, blayout>::Profile) \
.def("Run", &CKGemmFastGelu<type, alayout, blayout>::Run) \
.def("ListOps", &CKGemmFastGelu<type, alayout, blayout>::ListOps) \
.def("SelectOp", &CKGemmFastGelu<type, alayout, blayout>::SelectOp);
#define REGISTER_OP(type, opa, opb, layout_string) \
py::class_<CKGemmFastGelu<type, opa, opb>>(m, "CKGemmFastGelu_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &CKGemmFastGelu<type, opa, opb>::SetRepeats) \
.def("Profile", &CKGemmFastGelu<type, opa, opb>::Profile) \
.def("Run", &CKGemmFastGelu<type, opa, opb>::Run) \
.def("ListOps", &CKGemmFastGelu<type, opa, opb>::ListOps) \
.def("SelectOp", &CKGemmFastGelu<type, opa, opb>::SelectOp);
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, Row, Row, "NN"); \
REGISTER_OP(type, Row, Col, "NT"); \
REGISTER_OP(type, Col, Row, "TN"); \
REGISTER_OP(type, Col, Col, "TT");
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_OP(type, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_OP(type, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_OP(type, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_OP_FOR_ALL_TRANSAB(float);

42
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_hipblaslt.cu
View file

@ -23,7 +23,7 @@ namespace onnxruntime {
using namespace rocm::tunable::blas::internal;
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmFastGeluHipBlasLt : public IKernelExplorer {
public:
GemmFastGeluHipBlasLt(BlasOp opa, BlasOp opb,
@ -53,7 +53,7 @@ class GemmFastGeluHipBlasLt : public IKernelExplorer {
params_.c = static_cast<T*>(c.ptr());
params_.ldc = ldc;
for (auto&& [type_string, op] : GetHipBlasLtGemmFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetHipBlasLtGemmFastGeluTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -89,26 +89,26 @@ class GemmFastGeluHipBlasLt : public IKernelExplorer {
size_t selected_op_{};
};
#define REGISTER_OP(type, alayout, blayout, layout_string) \
py::class_<GemmFastGeluHipBlasLt<type, alayout, blayout>>(m, "GemmFastGeluHipBlasLt_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &GemmFastGeluHipBlasLt<type, alayout, blayout>::SetRepeats) \
.def("Profile", &GemmFastGeluHipBlasLt<type, alayout, blayout>::Profile) \
.def("Run", &GemmFastGeluHipBlasLt<type, alayout, blayout>::Run) \
.def("ListOps", &GemmFastGeluHipBlasLt<type, alayout, blayout>::ListOps) \
.def("SelectOp", &GemmFastGeluHipBlasLt<type, alayout, blayout>::SelectOp);
#define REGISTER_OP(type, opa, opb, layout_string) \
py::class_<GemmFastGeluHipBlasLt<type, opa, opb>>(m, "GemmFastGeluHipBlasLt_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &GemmFastGeluHipBlasLt<type, opa, opb>::SetRepeats) \
.def("Profile", &GemmFastGeluHipBlasLt<type, opa, opb>::Profile) \
.def("Run", &GemmFastGeluHipBlasLt<type, opa, opb>::Run) \
.def("ListOps", &GemmFastGeluHipBlasLt<type, opa, opb>::ListOps) \
.def("SelectOp", &GemmFastGeluHipBlasLt<type, opa, opb>::SelectOp);
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, Row, Row, "NN"); \
REGISTER_OP(type, Row, Col, "NT"); \
REGISTER_OP(type, Col, Row, "TN"); \
REGISTER_OP(type, Col, Col, "TT");
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_OP(type, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_OP(type, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_OP(type, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_OP_FOR_ALL_TRANSAB(float);

42
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu
View file

@ -17,7 +17,7 @@ using namespace onnxruntime::contrib::rocm::blas::internal;
namespace py = pybind11;
namespace onnxruntime {
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmFastGeluTunable : public IKernelExplorer {
public:
GemmFastGeluTunable(BlasOp opa, BlasOp opb,
@ -72,29 +72,29 @@ class GemmFastGeluTunable : public IKernelExplorer {
using ParamsT = GemmFastGeluParams<T>;
ParamsT params_{};
rocblas_handle rocblas_handle_;
GemmFastGeluTunableOp<T, ALayout, BLayout> op_{};
GemmFastGeluTunableOp<T, OpA, OpB> op_{};
};
#define REGISTER_OP(type, alayout, blayout, layout_string) \
py::class_<GemmFastGeluTunable<type, alayout, blayout>>(m, "GemmFastGeluTunable_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &GemmFastGeluTunable<type, alayout, blayout>::SetRepeats) \
.def("Profile", &GemmFastGeluTunable<type, alayout, blayout>::Profile) \
.def("Run", &GemmFastGeluTunable<type, alayout, blayout>::Run) \
.def("ListOps", &GemmFastGeluTunable<type, alayout, blayout>::ListOps) \
.def("SelectOp", &GemmFastGeluTunable<type, alayout, blayout>::SelectOp);
#define REGISTER_OP(type, opa, opb, layout_string) \
py::class_<GemmFastGeluTunable<type, opa, opb>>(m, "GemmFastGeluTunable_" #type "_" layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &GemmFastGeluTunable<type, opa, opb>::SetRepeats) \
.def("Profile", &GemmFastGeluTunable<type, opa, opb>::Profile) \
.def("Run", &GemmFastGeluTunable<type, opa, opb>::Run) \
.def("ListOps", &GemmFastGeluTunable<type, opa, opb>::ListOps) \
.def("SelectOp", &GemmFastGeluTunable<type, opa, opb>::SelectOp);
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, Row, Row, "NN"); \
REGISTER_OP(type, Row, Col, "NT"); \
REGISTER_OP(type, Col, Row, "TN"); \
REGISTER_OP(type, Col, Col, "TT");
#define REGISTER_OP_FOR_ALL_TRANSAB(type) \
REGISTER_OP(type, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_OP(type, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_OP(type, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_OP(type, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_OP_FOR_ALL_TRANSAB(float);

72
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_hipblaslt.cu
View file

@ -25,7 +25,7 @@ namespace onnxruntime {
using namespace rocm::tunable::blas::internal;
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmHipBlasLt : public IKernelExplorer {
public:
GemmHipBlasLt(BlasOp opa, BlasOp opb,
@ -54,7 +54,7 @@ class GemmHipBlasLt : public IKernelExplorer {
params_.c = static_cast<T*>(c.ptr());
params_.ldc = ldc;
for (auto&& [type_string, op] : GetHipBlasLtGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetHipBlasLtGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -90,7 +90,7 @@ class GemmHipBlasLt : public IKernelExplorer {
size_t selected_op_{};
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class StridedBatchedGemmHipBlasLt : public IKernelExplorer {
public:
StridedBatchedGemmHipBlasLt(
@ -125,7 +125,7 @@ class StridedBatchedGemmHipBlasLt : public IKernelExplorer {
params_.stride_c = stride_c;
params_.batch = batch;
for (auto&& [type_string, op] : GetHipBlasLtStridedBatchedGemmTypeStringAndOps<T, ALayout, BLayout>()) {
for (auto&& [type_string, op] : GetHipBlasLtStridedBatchedGemmTypeStringAndOps<T, OpA, OpB>()) {
type_strings_.emplace_back(std::move(type_string));
ops_.emplace_back(std::move(op));
}
@ -161,44 +161,44 @@ class StridedBatchedGemmHipBlasLt : public IKernelExplorer {
size_t selected_op_{};
};
#define REGISTER_OP_COMMON(type, dtype, alayout, blayout, layout_string) \
py::class_<type<dtype, alayout, blayout>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, alayout, blayout>::SetRepeats) \
.def("Profile", &type<dtype, alayout, blayout>::Profile) \
.def("Run", &type<dtype, alayout, blayout>::Run) \
.def("ListOps", &type<dtype, alayout, blayout>::ListOps) \
.def("SelectOp", &type<dtype, alayout, blayout>::SelectOp)
#define REGISTER_OP_COMMON(type, dtype, opa, opb, layout_string) \
py::class_<type<dtype, opa, opb>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, opa, opb>::SetRepeats) \
.def("Profile", &type<dtype, opa, opb>::Profile) \
.def("Run", &type<dtype, opa, opb>::Run) \
.def("ListOps", &type<dtype, opa, opb>::ListOps) \
.def("SelectOp", &type<dtype, opa, opb>::SelectOp)
#define REGISTER_GEMM_HIPBLASLT(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(GemmHipBlasLt, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
#define REGISTER_GEMM_HIPBLASLT(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(GemmHipBlasLt, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
DeviceArray&, int64_t>());
#define REGISTER_GEMM_HIPBLASLT_FOR_ALL_TRANSAB(dtype) \
REGISTER_GEMM_HIPBLASLT(dtype, Row, Row, "NN"); \
REGISTER_GEMM_HIPBLASLT(dtype, Row, Col, "NT"); \
REGISTER_GEMM_HIPBLASLT(dtype, Col, Row, "TN"); \
REGISTER_GEMM_HIPBLASLT(dtype, Col, Col, "TT");
#define REGISTER_GEMM_HIPBLASLT_FOR_ALL_TRANSAB(dtype) \
REGISTER_GEMM_HIPBLASLT(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_GEMM_HIPBLASLT(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_GEMM_HIPBLASLT(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_GEMM_HIPBLASLT(dtype, BlasOp::T, BlasOp::T, "TT");
#define REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(StridedBatchedGemmHipBlasLt, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
#define REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(StridedBatchedGemmHipBlasLt, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
int64_t>());
#define REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT_FOR_ALL_TRANSAB(dtype) \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, Row, Row, "NN"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, Row, Col, "NT"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, Col, Row, "TN"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, Col, Col, "TT");
#define REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT_FOR_ALL_TRANSAB(dtype) \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_STRIDEDBATCHEDGEMM_HIPBLASLT(dtype, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_GEMM_HIPBLASLT_FOR_ALL_TRANSAB(float);

102
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu
View file

@ -19,7 +19,7 @@ using namespace onnxruntime::rocm::tunable::blas::internal;
namespace onnxruntime {
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class GemmTunable : public IKernelExplorer {
public:
GemmTunable(BlasOp opa, BlasOp opb,
@ -73,11 +73,11 @@ class GemmTunable : public IKernelExplorer {
ParamsT params_;
// tunable is stateful, store it as an instance
GemmTunableOp<T, ALayout, BLayout> op_{};
GemmTunableOp<T, OpA, OpB> op_{};
rocblas_handle rocblas_handle_;
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class BatchedGemmTunable : public IBatchedGemmKernelExplorer<T> {
public:
BatchedGemmTunable(BlasOp opa, BlasOp opb,
@ -135,11 +135,11 @@ class BatchedGemmTunable : public IBatchedGemmKernelExplorer<T> {
ParamsT params_;
// tunable is stateful, store it as an instance
BatchedGemmTunableOp<T, ALayout, BLayout> op_{};
BatchedGemmTunableOp<T, OpA, OpB> op_{};
rocblas_handle rocblas_handle_;
};
template <typename T, typename ALayout, typename BLayout>
template <typename T, BlasOp OpA, BlasOp OpB>
class StridedBatchedGemmTunable : public IKernelExplorer {
public:
StridedBatchedGemmTunable(BlasOp opa, BlasOp opb,
@ -198,64 +198,64 @@ class StridedBatchedGemmTunable : public IKernelExplorer {
ParamsT params_;
// tunable is stateful, store it as an instance
StridedBatchedGemmTunableOp<T, ALayout, BLayout> op_{};
StridedBatchedGemmTunableOp<T, OpA, OpB> op_{};
rocblas_handle rocblas_handle_;
};
#define REGISTER_OP_COMMON(type, dtype, alayout, blayout, layout_string) \
py::class_<type<dtype, alayout, blayout>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, alayout, blayout>::SetRepeats) \
.def("Profile", &type<dtype, alayout, blayout>::Profile) \
.def("Run", &type<dtype, alayout, blayout>::Run) \
.def("ListOps", &type<dtype, alayout, blayout>::ListOps) \
.def("SelectOp", &type<dtype, alayout, blayout>::SelectOp)
#define REGISTER_OP_COMMON(type, dtype, opa, opb, layout_string) \
py::class_<type<dtype, opa, opb>>(m, #type "_" #dtype "_" layout_string) \
.def("SetRepeats", &type<dtype, opa, opb>::SetRepeats) \
.def("Profile", &type<dtype, opa, opb>::Profile) \
.def("Run", &type<dtype, opa, opb>::Run) \
.def("ListOps", &type<dtype, opa, opb>::ListOps) \
.def("SelectOp", &type<dtype, opa, opb>::SelectOp)
#define REGISTER_GEMM(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(GemmTunable, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
#define REGISTER_GEMM(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(GemmTunable, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
double, \
DeviceArray&, int64_t>())
#define REGISTER_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_GEMM(dtype, Row, Row, "NN"); \
REGISTER_GEMM(dtype, Row, Col, "NT"); \
REGISTER_GEMM(dtype, Col, Row, "TN"); \
REGISTER_GEMM(dtype, Col, Col, "TT");
#define REGISTER_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_GEMM(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_GEMM(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_GEMM(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_GEMM(dtype, BlasOp::T, BlasOp::T, "TT");
#define REGISTER_BATCHED_GEMM(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(BatchedGemmTunable, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
std::vector<DeviceArray>&, int64_t, \
std::vector<DeviceArray>&, int64_t, \
double, \
std::vector<DeviceArray>&, int64_t, \
#define REGISTER_BATCHED_GEMM(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(BatchedGemmTunable, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
std::vector<DeviceArray>&, int64_t, \
std::vector<DeviceArray>&, int64_t, \
double, \
std::vector<DeviceArray>&, int64_t, \
int64_t>())
#define REGISTER_BATCHED_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_BATCHED_GEMM(dtype, Row, Row, "NN"); \
REGISTER_BATCHED_GEMM(dtype, Row, Col, "NT"); \
REGISTER_BATCHED_GEMM(dtype, Col, Row, "TN"); \
REGISTER_BATCHED_GEMM(dtype, Col, Col, "TT");
#define REGISTER_BATCHED_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_BATCHED_GEMM(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_BATCHED_GEMM(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_BATCHED_GEMM(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_BATCHED_GEMM(dtype, BlasOp::T, BlasOp::T, "TT");
#define REGISTER_STRIDED_BATCHED_GEMM(dtype, alayout, blayout, layout_string) \
REGISTER_OP_COMMON(StridedBatchedGemmTunable, dtype, alayout, blayout, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
#define REGISTER_STRIDED_BATCHED_GEMM(dtype, opa, opb, layout_string) \
REGISTER_OP_COMMON(StridedBatchedGemmTunable, dtype, opa, opb, layout_string) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
DeviceArray&, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, int64_t, \
int64_t>())
#define REGISTER_STRIDED_BATCHED_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_STRIDED_BATCHED_GEMM(dtype, Row, Row, "NN"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, Row, Col, "NT"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, Col, Row, "TN"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, Col, Col, "TT");
#define REGISTER_STRIDED_BATCHED_GEMM_FOR_ALL_TRANSAB(dtype) \
REGISTER_STRIDED_BATCHED_GEMM(dtype, BlasOp::N, BlasOp::N, "NN"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, BlasOp::N, BlasOp::T, "NT"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, BlasOp::T, BlasOp::N, "TN"); \
REGISTER_STRIDED_BATCHED_GEMM(dtype, BlasOp::T, BlasOp::T, "TT");
KE_REGISTER(m) {
REGISTER_GEMM_FOR_ALL_TRANSAB(float);