junrushao1994 commented on a change in pull request #9871:
URL: https://github.com/apache/tvm/pull/9871#discussion_r780587954
##########
File path: include/tvm/tir/function.h
##########
@@ -187,6 +187,56 @@ class LinkedParam : public ObjectRef {
TVM_DEFINE_OBJECT_REF_COW_METHOD(LinkedParamNode);
};
+/*!
+ * \brief Tensor intrinsics for tensorization
+ */
+class TensorIntrinNode : public Object {
+ public:
+ /*! \brief The function to describe the computation. */
+ PrimFunc description;
+ /*! \brief The intrinsic function for lower-level implementation. */
+ PrimFunc implementation;
+
+ void VisitAttrs(AttrVisitor* v) {
+ v->Visit("description", &description);
+ v->Visit("implementation", &implementation);
+ }
+
+ static constexpr const char* _type_key = "tir.TensorIntrin";
+ TVM_DECLARE_FINAL_OBJECT_INFO(TensorIntrinNode, Object);
+};
+
+/*!
+ * \brief Managed reference to TensorIntrinNode.
+ */
+class TensorIntrin : public ObjectRef {
+ public:
+ /*!
+ * \brief Constructor
+ * \param desc_func The function to describe the computation.
+ * \param intrin_func The intrinsic function for lower-level implementation.
+ */
+ TVM_DLL explicit TensorIntrin(PrimFunc desc_func, PrimFunc intrin_func);
+
+ /*!
+ * \brief Create and register a TensorIntrin. After registration, the
TensorIntrin can be looked
+ * up with its name.
+ * \param name The name of the TensorIntrin to register
+ * \param desc_func The function to describe the computation.
+ * \param intrin_func The intrinsic function for lower-level implementation.
+ * \return The created TensorIntrin.
+ */
+ TVM_DLL static TensorIntrin Register(String name, PrimFunc desc_func,
PrimFunc intrin_func);
Review comment:
Let's be more idiomatic in registry design
```suggestion
TVM_DLL static TensorIntrin Register(String name, TensorIntrin intrin);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
