fixup! [TIR] Add schedule primitive TransformBlockLayout

Fix doc

Author: vinx13

include/tvm/tir/schedule/schedule.h

@@ -550,9 +550,8 @@ class ScheduleNode : public runtime::Object {
    * \details The block iters and the block body are transformed by the given index_map.
    * Outer loops corresponding to each new block iter are regenerated.
    * The index_map is required to be bijective affine since we need its inverse mapping.
-   * \param self The state of the schedule
-   * \param block_sref The block sref that refers to the block to be transformed
-   * \param affine_index_map The transformation to apply.
+   * \param block_rv The block to be transformed
+   * \param index_map The transformation to apply.
    */
   virtual void TransformBlockLayout(const BlockRV& block_rv, const IndexMap& index_map) = 0;
 

python/tvm/tir/schedule/schedule.py

@@ -2296,7 +2296,7 @@ def transform_block_layout(
 
         Parameters
         ----------
-        block_rv : BlockRV
+        block : BlockRV
             The block to be transformed
 
         index_map : Union[IndexMap, Callable]

src/tir/schedule/primitive.h

@@ -449,7 +449,7 @@ TVM_DLL void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int
  * The index_map is required to be bijective affine since we need its inverse mapping.
  * \param self The state of the schedule
  * \param block_sref The block sref that refers to the block to be transformed
- * \param affine_index_map The transformation to apply.
+ * \param index_map The transformation to apply.
  */
 TVM_DLL void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref,
                                   const IndexMap& index_map);

src/tir/schedule/primitive/layout_transformation.cc

@@ -405,7 +405,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref,
   Array<IterVar> new_block_iters;  // new block iters
   Array<PrimExpr> new_block_vars;  // iter_var->var of new block iters
   for (size_t i = 0; i < index_map->final_indices.size(); ++i) {
-    Var new_block_var{"bv" + std::to_string(i), DataType::Int(32)};
+    Var new_block_var{"v" + std::to_string(i), DataType::Int(32)};
     new_block_vars.push_back(new_block_var);
     IterVarType iter_type = DetectNewBlockIterType(transformed_block_iters[i], block_iter_type);
     if (iter_type == kOpaque) {