Add_SimpleCategorical

cinn/backends/ir_schedule_test.cc

@@ -2921,5 +2921,23 @@ TEST(IrSchedule, GetChildBlocks) {
   ASSERT_EQ(utils::GetStreamCnt(ir_sch.GetChildBlocks(root_block)), expected_expr);
 }
 
+TEST(IrSchedule, SimpleCategorical) {
+  Context::Global().ResetNameId();
+  Expr M(32);
+  Expr N(32);
+  Placeholder<int> A("A", {M, N});
+  auto B = Compute(
+      {M, N}, [&](Var i, Var j) { return A(i, j); }, "B");
+  poly::StageMap stages = CreateStages({A, B});
+
+  auto funcs = cinn::lang::LowerVec(
+      "test_simplecategorical", stages, {A, B}, {}, {}, nullptr, common::DefaultHostTarget(), true);
+
+  ir::IRSchedule ir_sch(ir::ModuleExpr({funcs[0]->body}));
+  Expr result = ir_sch.SimpleCategorical({1, 2, 3}, {1.0, 2.0, 3.0});
+  LOG(INFO) << "SimpleCategorical result: " << result;
+  ASSERT_EQ(result.type(), Int(32));
+}
+
 }  // namespace backends
 }  // namespace cinn

cinn/ir/ir_schedule.cc

@@ -102,6 +102,7 @@ class ScheduleImpl {
   void FlattenLoops(const std::vector<Expr>& loops, const bool force_flat = false);
   void CopyTransformAndLoopInfo(const Expr& block, const Expr& block_target);
   void CopyTransformAndLoopInfo(const std::string& block_name, const std::string& block_target_name);
+  Expr SimpleCategorical(const std::vector<int>& candidates, const std::vector<float>& probs);
 
  private:
   void Replace(const Expr& src_sref, const Expr& tgt_stmt);
@@ -1692,6 +1693,25 @@ void ScheduleImpl::FlattenLoops(const std::vector<Expr>& loops, const bool flat_
   this->Replace(loops[0], loop);
 }
 
+Expr ScheduleImpl::SimpleCategorical(const std::vector<int>& candidates, const std::vector<float>& probs) {
+  int i = -1;
+  // check two sizes
+  CHECK_EQ(candidates.size(), probs.size()) << "candidates and probs must have same size.";
+  // generate
+  std::vector<double> weights;
+  for (auto p : probs) {
+    weights.push_back(p);
+  }
+  std::discrete_distribution<int32_t> prob_int(weights.begin(), weights.end());
+
+  std::random_device seed;
+  std::default_random_engine engine(seed());
+  i           = prob_int(engine);
+  auto result = candidates[i];
+  Expr result_Expr(result);
+  return result_Expr;
+}
+
 void ScheduleImpl::CopyTransformAndLoopInfo(const std::string& block_name, const std::string& block_target_name) {
   auto block        = this->GetBlock(block_name);
   auto block_target = this->GetBlock(block_target_name);
@@ -2093,6 +2113,10 @@ std::vector<Expr> IRSchedule::SamplePerfectTile(const Expr& loop, int n, int max
                                    {result}));
   return result;
 }
-
+Expr IRSchedule::SimpleCategorical(const std::vector<int>& candidates, const std::vector<float>& probs) {
+  auto result = impl_->SimpleCategorical(candidates, probs);
+  trace_.Append(ScheduleDesc::Step("SimpleCategorical", {}, {{"candidates", candidates}, {"probs", probs}}, {result}));
+  return result;
+}
 }  // namespace ir
 }  // namespace cinn

cinn/ir/ir_schedule.h

@@ -373,6 +373,14 @@ class IRSchedule {
    */
   std::vector<Expr> SamplePerfectTile(const Expr& loop, int n, int max_innermost_factor);
 
+  /**
+   * \brief Randomly sample an integer according to the given distribution.
+   * @param candidates Candidate set of integers.
+   * @param probs Probability distribution of candidate integer set.
+   * @return Random variables sampled.
+   */
+  Expr SimpleCategorical(const std::vector<int>& candidates, const std::vector<float>& probs);
+
  private:
   std::unique_ptr<ScheduleImpl> impl_;
   mutable ScheduleDesc trace_;  // trace the scheduling process