add additional temp space

apache · Aug 20, 2019 · e658902 · e658902
1 parent a1929d4
commit e658902
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/operator/tensor/ordering_op-inl.h b/src/operator/tensor/ordering_op-inl.h
@@ -419,6 +419,11 @@ void TopKImpl(const RunContext &ctx,
       mxnet::op::SortByKeyWorkspaceSize<index_t, DType, xpu>(src.Size()),
       mxnet::op::SortByKeyWorkspaceSize<DType, index_t, xpu>(src.Size()));
 
+  // Additional temp space for gpu full sorts for batch ids.
+  temp_size += PadBytes(sizeof(index_t) * src.Size(), alignment);
+  // Temp space for cpu sorts.
+  temp_size = std::max(temp_size, sizeof(DType) * src.Size());
+
   size_t workspace_size = temp_size + PadBytes(sizeof(DType) * src.Size(), alignment)
                                     + PadBytes(sizeof(index_t) * src.Size(), alignment);
   if (param.ret_typ == topk_enum::kReturnMask) {