improve sampling time

Differential Revision: D60742125

Pull Request resolved: #4644

Author: cccclai

extension/llm/sampler/sampler.cpp

@@ -33,6 +33,7 @@
  */
 
 #include <executorch/extension/llm/sampler/sampler.h>
+#include <algorithm>
 
 namespace torch {
 namespace executor {
@@ -66,18 +67,6 @@ int32_t Sampler::sample_mult(T* probabilities, float coin) {
   return vocab_size_ - 1; // in case of rounding errors
 }
 
-template <typename T>
-static int32_t compare(const void* a, const void* b) {
-  ProbIndex<T>* a_ = (ProbIndex<T>*)a;
-  ProbIndex<T>* b_ = (ProbIndex<T>*)b;
-  if (a_->prob > b_->prob) {
-    return -1;
-  } else if (a_->prob < b_->prob) {
-    return 1;
-  }
-  return 0;
-}
-
 template <typename T>
 int32_t Sampler::sample_topp(T* probabilities, float coin) {
   // top-p sampling (or "nucleus sampling") samples from the smallest set of
@@ -100,7 +89,11 @@ int32_t Sampler::sample_topp(T* probabilities, float coin) {
       n0++;
     }
   }
-  qsort(probindex.get(), n0, sizeof(ProbIndex<T>), compare<T>);
+
+  auto compare = [](const ProbIndex<T>& a, const ProbIndex<T>& b) {
+    return a.prob > b.prob;
+  };
+  std::sort(probindex.get(), probindex.get() + n0, compare);
 
   // truncate the list where cumulative probability exceeds topp
   T cumulative_prob = 0;