add doc

Author: masahi

python/tvm/meta_schedule/schedule_rule/auto_tensorize.py

@@ -15,9 +15,28 @@
 # specific language governing permissions and limitations
 # under the License.
 """Auto tensorization related components."""
+from typing import Optional
 
 from .. import _ffi_api
+from tvm.tir.schedule import Schedule, BlockRV, LoopRV
 
 
-def tiling_with_tensor_intrin(sch, block, intrin_name):
-    return _ffi_api.TilingwithTensorIntrin(sch, block, intrin_name)
+def tiling_with_tensor_intrin(sch: Schedule, block: BlockRV, intrin_name: str) -> Optional[LoopRV]:
+    """Tile a subset of loops in the block according to the given tensor intrinsic.
+
+    Parameters
+    ----------
+    sch : Schedule
+        The schedule to which tiling is applied
+    block : BlockRV
+        The block whose subset of loops will be tiled
+    intrin_name : str
+        The name of a tensor intrinsic, must be registerd via TensorIntrin.register(...) beforehand
+
+    Returns
+    -------
+    tiled_loop_rv : Optional[LoopRV]
+        LoopRV corresponding to the outermost loop of a block tiled according to the given intrin
+        NullOpt if no valid loop mapping is found
+    """
+    return _ffi_api.TilingwithTensorIntrin(sch, block, intrin_name)  # type: ignore

src/meta_schedule/schedule_rule/auto_tensorize.h

@@ -24,6 +24,16 @@
 namespace tvm {
 namespace meta_schedule {
 
+/*!
+ * \brief Tile a subset of loops in the block according to the given tensor intrinsic.
+ * \param self The schedule to which tiling is applied
+ * \param block_rv The block whose subset of loops will be tiled
+ * \param intrin_name The name of a tensor intrinsic, must be registerd via
+ * TensorIntrin.register(...) beforehand
+ * \return LoopRV corresponding to the outermost loop of a
+ * block tiled according to the given intrin, NullOpt if a valid loop mapping is not found
+ */
+
 Optional<tir::LoopRV> TilingwithTensorIntrin(const tir::Schedule& sch, const tir::BlockRV& block_rv,
                                              const String& intrin_name);