[BugFix] Fix binary search & SpIterVar (#7)

yzh119 · Feb 15, 2022 · 99f6a35 · 99f6a35
1 parent 88451ba
commit 99f6a35
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/include/tvm/tir/op.h b/include/tvm/tir/op.h
@@ -793,6 +793,30 @@ TVM_DLL PrimExpr round(PrimExpr x, Span span = Span());
  */
 TVM_DLL PrimExpr nearbyint(PrimExpr x, Span span = Span());
 
+/*!
+ * \brief Lower bound function for binary search
+ * \param arr The buffer variable of the array to be looked up in
+ * \param val The value to be looked up in the array
+ * \param l The left boundary of the look-up range (inclusive)
+ * \param r The right boundary of the look-up range (exclusive)
+ * \param span The location of this operation in the source
+ * \return The look-up result
+ */
+TVM_DLL PrimExpr lower_bound(tir::Var arr, PrimExpr val, PrimExpr l, PrimExpr r,
+                             Span span = Span());
+
+/*!
+ * \brief Upper bound function for binary search
+ * \param arr The buffer variable of the array to be looked up in
+ * \param val The value to be looked up in the array
+ * \param l The left boundary of the look-up range (inclusive)
+ * \param r The right boundary of the look-up range (exclusive)
+ * \param span The location of this operation in the source
+ * \return The look-up result
+ */
+TVM_DLL PrimExpr upper_bound(tir::Var arr, PrimExpr val, PrimExpr l, PrimExpr r,
+                             Span span = Span());
+
 /*!
  * \brief Calculate trunc(x)
  * \param x The input expression.

diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
@@ -805,12 +805,12 @@ PrimExpr nearbyint(PrimExpr x, Span span) {
 TIR_REGISTER_PURE_UNARY_OP("tir.nearbyint");
 
 // lower_bound
-PrimExpr lower_bound(Var arr, PrimExpr val, PrimExpr l, PrimExpr r, Span span) {
+PrimExpr lower_bound(tir::Var arr, PrimExpr val, PrimExpr l, PrimExpr r, Span span) {
   return tir::Call({kDLInt, 32, 1}, builtin::tvm_lower_bound(), {arr, val, l, r}, span);
 }
 
 // upper_bound
-PrimExpr upper_bound(Var arr, PrimExpr val, PrimExpr l, PrimExpr r, Span span) {
+PrimExpr upper_bound(tir::Var arr, PrimExpr val, PrimExpr l, PrimExpr r, Span span) {
   return tir::Call({kDLInt, 32, 1}, builtin::tvm_upper_bound(), {arr, val, l, r}, span);
 }