[AMD] Skip unused TOPK v2 plan kernel on ROCm (#37580)
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
This commit is contained in:
@@ -467,8 +467,13 @@ struct TopKKernel {
|
|||||||
.with_device(device_)
|
.with_device(device_)
|
||||||
.verify(metadata);
|
.verify(metadata);
|
||||||
|
|
||||||
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
|
||||||
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
|
RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
|
||||||
|
#ifdef USE_ROCM
|
||||||
|
// ROCm compiles out the cluster path, the only consumer of this plan.
|
||||||
|
(void)static_cluster_threshold;
|
||||||
|
return;
|
||||||
|
#else
|
||||||
|
const auto batch_size = static_cast<uint32_t>(B.unwrap());
|
||||||
const auto device = device_.unwrap();
|
const auto device = device_.unwrap();
|
||||||
LaunchKernel(1, kBlockSize, device)( //
|
LaunchKernel(1, kBlockSize, device)( //
|
||||||
topk_plan,
|
topk_plan,
|
||||||
@@ -476,6 +481,7 @@ struct TopKKernel {
|
|||||||
static_cast<PlanItem*>(metadata.data_ptr()),
|
static_cast<PlanItem*>(metadata.data_ptr()),
|
||||||
batch_size,
|
batch_size,
|
||||||
static_cluster_threshold);
|
static_cluster_threshold);
|
||||||
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
static void transform_paged(
|
static void transform_paged(
|
||||||
|
|||||||
Reference in New Issue
Block a user