update postproc.h

Author: masahi

include/tvm/meta_schedule/postproc.h

@@ -149,10 +149,11 @@ class Postproc : public runtime::ObjectRef {
    */
   TVM_DLL static Postproc RewriteUnboundBlock(int max_threadblock);
   /*!
-   * \brief Create a postprocessor that tensorize Tensor Core related components
+   * \brief Create a postprocessor that applies tensorization to annotated blocks
+   * \param Whether or not vectorize the initialization loop produced by DecomposeReduction
    * \return The postprocessor created.
    */
-  TVM_DLL static Postproc RewriteTensorCore();
+  TVM_DLL static Postproc RewriteTensorize(bool vectorize_init_loop = false);
 
   /*!
    * \brief Creates a postprocessor that verifies if the GPU code is correct

python/tvm/meta_schedule/postproc/rewrite_tensorize.py

@@ -28,7 +28,7 @@ class RewriteTensorize(Postproc):
     Parameters
     ----------
     vectorize_init_loop : bool
-       Whether or not vectorize the initialization loop produced by decompose_reduction
+       Whether or not vectorize the initialization loop produced by DecomposeReduction
 
     """
 

src/meta_schedule/postproc/rewrite_tensorize.cc

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-#include <tvm/runtime/container/base.h>
+#include <tvm/meta_schedule/postproc.h>
 
 #include <algorithm>
 
@@ -91,14 +91,15 @@ bool RewriteTensorizeNode::Apply(const tir::Schedule& sch) {
   return true;
 }
 
-Postproc RewriteTensorize(bool vectorize_init_loop) {
+Postproc Postproc::RewriteTensorize(bool vectorize_init_loop) {
   ObjectPtr<RewriteTensorizeNode> n = make_object<RewriteTensorizeNode>();
   n->vectorize_init_loop = vectorize_init_loop;
   return Postproc(n);
 }
 
 TVM_REGISTER_NODE_TYPE(RewriteTensorizeNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteTensorize").set_body_typed(RewriteTensorize);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteTensorize")
+    .set_body_typed(Postproc::RewriteTensorize);
 
 }  // namespace meta_schedule
 }  // namespace tvm