From 4d09227134b208e65b7c6d5f6b47e7b6c7a3c3ae Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 13 May 2026 14:40:38 +0800 Subject: [PATCH 1/2] feat: add `searchsorted` base --- src/base/searchsorted.h | 76 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 src/base/searchsorted.h diff --git a/src/base/searchsorted.h b/src/base/searchsorted.h new file mode 100644 index 000000000..b45af2688 --- /dev/null +++ b/src/base/searchsorted.h @@ -0,0 +1,76 @@ +#ifndef INFINI_OPS_BASE_SEARCHSORTED_H_ +#define INFINI_OPS_BASE_SEARCHSORTED_H_ + +#include "operator.h" + +namespace infini::ops { + +class Searchsorted : public Operator { + public: + Searchsorted(const Tensor sorted_sequence, const Tensor input, + const bool out_int32, const bool right, Tensor out) + : sorted_sequence_shape_{sorted_sequence.shape()}, + sorted_sequence_strides_{sorted_sequence.strides()}, + sorted_sequence_type_{sorted_sequence.dtype()}, + input_shape_{input.shape()}, + input_strides_{input.strides()}, + input_type_{input.dtype()}, + out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()}, + out_int32_{out_int32}, + right_{right}, + device_index_{out.device().index()} {} + + Searchsorted(const Tensor sorted_sequence, const double input, + const bool out_int32, const bool right, Tensor out) + : sorted_sequence_shape_{sorted_sequence.shape()}, + sorted_sequence_strides_{sorted_sequence.strides()}, + sorted_sequence_type_{sorted_sequence.dtype()}, + out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()}, + out_int32_{out_int32}, + right_{right}, + input_{input}, + device_index_{out.device().index()} {} + + virtual void operator()(const Tensor sorted_sequence, const Tensor input, + const bool out_int32, const bool right, + Tensor out) const = 0; + + virtual void operator()(const Tensor sorted_sequence, const double input, + const bool out_int32, const bool right, + Tensor out) const = 0; + + protected: + Tensor::Shape sorted_sequence_shape_; + + Tensor::Strides sorted_sequence_strides_; + + DataType sorted_sequence_type_; + + Tensor::Shape input_shape_; + + Tensor::Strides input_strides_; + + DataType input_type_; + + Tensor::Shape out_shape_; + + Tensor::Strides out_strides_; + + DataType out_type_; + + bool out_int32_{}; + + bool right_{}; + + double input_{}; + + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif \ No newline at end of file From e50d7bb559cd952240580d1f9659ed322f644a59 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 18 May 2026 19:27:30 +0800 Subject: [PATCH 2/2] fix: add searchsorted side and sorter args --- src/base/searchsorted.h | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/src/base/searchsorted.h b/src/base/searchsorted.h index b45af2688..91f6bee0b 100644 --- a/src/base/searchsorted.h +++ b/src/base/searchsorted.h @@ -1,6 +1,8 @@ #ifndef INFINI_OPS_BASE_SEARCHSORTED_H_ #define INFINI_OPS_BASE_SEARCHSORTED_H_ +#include + #include "operator.h" namespace infini::ops { @@ -8,7 +10,9 @@ namespace infini::ops { class Searchsorted : public Operator { public: Searchsorted(const Tensor sorted_sequence, const Tensor input, - const bool out_int32, const bool right, Tensor out) + const bool out_int32, const bool right, + std::optional side, std::optional sorter, + Tensor out) : sorted_sequence_shape_{sorted_sequence.shape()}, sorted_sequence_strides_{sorted_sequence.strides()}, sorted_sequence_type_{sorted_sequence.dtype()}, @@ -20,10 +24,14 @@ class Searchsorted : public Operator { out_type_{out.dtype()}, out_int32_{out_int32}, right_{right}, + side_{side}, + has_sorter_{sorter.has_value()}, device_index_{out.device().index()} {} Searchsorted(const Tensor sorted_sequence, const double input, - const bool out_int32, const bool right, Tensor out) + const bool out_int32, const bool right, + std::optional side, std::optional sorter, + Tensor out) : sorted_sequence_shape_{sorted_sequence.shape()}, sorted_sequence_strides_{sorted_sequence.strides()}, sorted_sequence_type_{sorted_sequence.dtype()}, @@ -32,16 +40,20 @@ class Searchsorted : public Operator { out_type_{out.dtype()}, out_int32_{out_int32}, right_{right}, + side_{side}, + has_sorter_{sorter.has_value()}, input_{input}, device_index_{out.device().index()} {} virtual void operator()(const Tensor sorted_sequence, const Tensor input, const bool out_int32, const bool right, - Tensor out) const = 0; + std::optional side, + std::optional sorter, Tensor out) const = 0; virtual void operator()(const Tensor sorted_sequence, const double input, const bool out_int32, const bool right, - Tensor out) const = 0; + std::optional side, + std::optional sorter, Tensor out) const = 0; protected: Tensor::Shape sorted_sequence_shape_; @@ -66,6 +78,10 @@ class Searchsorted : public Operator { bool right_{}; + std::optional side_{}; + + bool has_sorter_{false}; + double input_{}; int device_index_{0}; @@ -73,4 +89,4 @@ class Searchsorted : public Operator { } // namespace infini::ops -#endif \ No newline at end of file +#endif