|
|
|
@@ -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());
|
|
|
|
|