CUTLASS 3.0.0 (#786)

* CUTLASS 3.0.0
This commit is contained in:
Vijay Thakkar
2023-01-23 20:55:28 -05:00
committed by GitHub
parent 66d9cddc83
commit 277bd6e537
377 changed files with 76396 additions and 1186 deletions
@@ -537,6 +537,13 @@ Status GemmOperationProfiler::initialize_workspace(
gemm_workspace_.Reference->copy_from_device(gemm_workspace_.C->data());
// NOTE: the leading non-batch strides are duplicated here for 3.0 API kernels
gemm_workspace_.arguments.problem_size = {int(problem_.m), int(problem_.n), int(problem_.k)};
gemm_workspace_.arguments.batch_count = problem_.batch_count;
gemm_workspace_.arguments.lda = problem_.lda;
gemm_workspace_.arguments.ldb = problem_.ldb;
gemm_workspace_.arguments.ldc = problem_.ldc;
gemm_workspace_.arguments.ldd = problem_.ldc;
gemm_workspace_.arguments.batch_stride_A = gemm_workspace_.A->batch_stride();
gemm_workspace_.arguments.batch_stride_B = gemm_workspace_.B->batch_stride();
gemm_workspace_.arguments.batch_stride_C = gemm_workspace_.C->batch_stride();
+24
View File
@@ -75,6 +75,9 @@ OperationProfiler::OperationProfiler(
{ArgumentTypeID::kInteger, {"cta_m", "threadblock-shape::m"}, "Threadblock shape in the M dimension"},
{ArgumentTypeID::kInteger, {"cta_n", "threadblock-shape::n"}, "Threadblock shape in the N dimension"},
{ArgumentTypeID::kInteger, {"cta_k", "threadblock-shape::k"}, "Threadblock shape in the K dimension"},
{ArgumentTypeID::kInteger, {"cluster_m", "cluster-shape::m"}, "Cluster shape in the M dimension"},
{ArgumentTypeID::kInteger, {"cluster_n", "cluster-shape::n"}, "Cluster shape in the N dimension"},
{ArgumentTypeID::kInteger, {"cluster_k", "cluster-shape::k"}, "Cluster shape in the K dimension"},
{ArgumentTypeID::kInteger, {"stages", "threadblock-stages"}, "Number of stages of threadblock-scoped matrix multiply"},
{ArgumentTypeID::kInteger, {"warps_m", "warp-count::m"}, "Number of warps within threadblock along the M dimension"},
{ArgumentTypeID::kInteger, {"warps_n", "warp-count::n"}, "Number of warps within threadblock along the N dimension"},
@@ -198,6 +201,24 @@ bool OperationProfiler::satisfies(
}
}
if (arg_as_int(int_value, "cluster_m", problem_space, problem)) {
if (int64_t(op_desc.tile_description.cluster_shape.m()) != int_value) {
return false;
}
}
if (arg_as_int(int_value, "cluster_n", problem_space, problem)) {
if (int64_t(op_desc.tile_description.cluster_shape.n()) != int_value) {
return false;
}
}
if (arg_as_int(int_value, "cluster_k", problem_space, problem)) {
if (int64_t(op_desc.tile_description.cluster_shape.k()) != int_value) {
return false;
}
}
if (arg_as_int(int_value, "stages", problem_space, problem)) {
if (int64_t(op_desc.tile_description.threadblock_stages) != int_value) {
return false;
@@ -596,6 +617,9 @@ void OperationProfiler::initialize_result_(
set_argument(result, "cta_m", problem_space, operation_desc.tile_description.threadblock_shape.m());
set_argument(result, "cta_n", problem_space, operation_desc.tile_description.threadblock_shape.n());
set_argument(result, "cta_k", problem_space, operation_desc.tile_description.threadblock_shape.k());
set_argument(result, "cluster_m", problem_space, operation_desc.tile_description.cluster_shape.m());
set_argument(result, "cluster_n", problem_space, operation_desc.tile_description.cluster_shape.n());
set_argument(result, "cluster_k", problem_space, operation_desc.tile_description.cluster_shape.k());
set_argument(result, "stages", problem_space, operation_desc.tile_description.threadblock_stages);
set_argument(result, "warps_m", problem_space, operation_desc.tile_description.warp_count.m());
set_argument(result, "warps_n", problem_space, operation_desc.tile_description.warp_count.n());