lint

Author: vinx13

python/tvm/tir/function.py

@@ -182,9 +182,7 @@ def __init__(self, desc_func, intrin_func):
 
     @staticmethod
     def register(name: str, desc_func: PrimFunc, intrin_func: PrimFunc):
-        return _ffi_api.TensorIntrinRegister(  # type: ignore
-            name, desc_func, intrin_func
-        )
+        return _ffi_api.TensorIntrinRegister(name, desc_func, intrin_func)  # type: ignore
 
     @staticmethod
     def get(name: str):

python/tvm/tir/schedule/schedule.py

@@ -1766,7 +1766,7 @@ def before_blockize(
                         T.reads(A[vi, vj])
                         T.writes(B[vi, vj])
                         B[vi, vj] = A[vi, vj] * T.float32(2)
-            
+
         Create the schedule and do set_scope:
 
         .. code-block:: python
@@ -1803,9 +1803,9 @@ def after_blockize(
         ----
         blockize requires there is exactly one block under the given loop and the bindings of the
         block are divisible by the subspace represented by the loops starting at the given loop.
-            
+
         """
-        
+
         return _ffi_api.ScheduleBlockize(self, loop)  # type: ignore # pylint: disable=no-member
 
     @type_checked
@@ -1860,7 +1860,7 @@ def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
                 A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
                 B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
                 C = T.match_buffer(c, (16, 16), align=128, offset_factor=1)
-            
+
                 with T.block("root"):
                     vi = T.axis.S(16, 0)
                     vj = T.axis.S(16, 0)
@@ -1871,14 +1871,14 @@ def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
                             vjj = T.axis.S(16, vj + j)
                             vkk = T.axis.R(16, vk + k)
                             C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]
-            
-            
+
+
             @T.prim_func
             def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
                 A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
                 B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
                 C = T.match_buffer(c, (16, 16), align=128, offset_factor=1)
-            
+
                 with T.block("root"):
                     vi = T.axis.S(16, 0)
                     vj = T.axis.S(16, 0)
@@ -1977,12 +1977,13 @@ def after_tensoirze(
                                     dtype="handle",
                                 )
                             )
-            
+
         """
         if isinstance(tensor_intrin, str):
             tensor_intrin = String(tensor_intrin)
         _ffi_api.ScheduleTensorize(  # type: ignore # pylint: disable=no-member
-            self, loop, tensor_intrin)
+            self, loop, tensor_intrin
+        )
 
     ########## Schedule: Annotation ##########
 

src/tir/ir/function.cc

@@ -80,8 +80,10 @@ TensorIntrin::TensorIntrin(PrimFunc desc_func, PrimFunc intrin_func) {
   CHECK_EQ(desc_func->buffer_map.size(), intrin_func->buffer_map.size());
 
   // check both functions' bodies are directly block
-  const auto* desc_realize = Downcast<BlockRealize>(desc_func->body)->block->body.as<BlockRealizeNode>();
-  const auto* intrin_realize = Downcast<BlockRealize>(intrin_func->body)->block->body.as<BlockRealizeNode>();
+  const auto* desc_realize =
+      Downcast<BlockRealize>(desc_func->body)->block->body.as<BlockRealizeNode>();
+  const auto* intrin_realize =
+      Downcast<BlockRealize>(intrin_func->body)->block->body.as<BlockRealizeNode>();
   CHECK(desc_realize != nullptr) << "description function's body expect a directly block";
   CHECK(intrin_realize != nullptr) << "intrinsic function's body expect a directly block";
 
@@ -144,7 +146,6 @@ TVM_REGISTER_GLOBAL("tir.PrimFunc")
       return PrimFunc(params, body, ret_type, buffer_map, attrs, span);
     });
 
-
 TVM_REGISTER_GLOBAL("tir.TensorIntrin")
     .set_body_typed([](PrimFunc desc_func, PrimFunc intrin_func) {
       return TensorIntrin(desc_func, intrin_func);

src/tir/schedule/analysis.h

@@ -447,9 +447,13 @@ bool CanReverseComputeAt(const ScheduleState& self, const StmtSRef& block_sref,
 using ExprComparator = ExprFunctor<bool(const PrimExpr& n, const PrimExpr& other)>;
 using StmtComparator = StmtFunctor<bool(const Stmt& n, const Stmt& other)>;
 
-/* \brief Deep comparison to check if two IR ASTs are equivalent */
+/*! \brief Deep comparison to check if two IR ASTs are equivalent */
 class TensorizeComparator : public ExprComparator, public StmtComparator {
  public:
+  /*!
+   * \brief Constructor of TensorizeComparator
+   * \param assert_mode Whether to raise an error if the two IR ASTs do not match.
+   */
   explicit TensorizeComparator(bool assert_mode = true) : assert_mode_(assert_mode) {}
 
   // Map from rhs buffer to lhs buffer
@@ -507,10 +511,10 @@ class TensorizeComparator : public ExprComparator, public StmtComparator {
 
  protected:
   bool assert_mode_;
-  bool is_scope_block = true, is_inner_block = true;
+  bool is_scope_block = true;
+  bool is_inner_block = true;
 };
 
-
 }  // namespace tir
 }  // namespace tvm
 

src/tir/schedule/analysis/analysis.cc

@@ -1374,7 +1374,7 @@ void CheckStorageScope(const ScheduleState& self, String storage_scope) {
   } catch (...) {
     throw InvalidStorageScopeError(self->mod, std::move(storage_scope));
   }
-};
+}
 
 /******** Tensorize Comparator ********/
 

src/tir/schedule/primitive/blockize_tensorize.cc

@@ -364,7 +364,7 @@ BufferRegion RelaxBlockizedInnerIters(const BufferRegion& buffer_region,
     new_region.push_back(relaxed_int_set[i].CoverRange(max_range));
   }
   return BufferRegion(buffer_region->buffer, std::move(new_region));
-};
+}
 
 /*!
  * \brief Generate the outer block after blockize.

src/tir/schedule/traced_schedule.cc

@@ -364,7 +364,8 @@ void TracedScheduleNode::Tensorize(const LoopRV& loop_rv, const String& intrin_n
 }
 
 void TracedScheduleNode::Tensorize(const LoopRV& loop_rv, const TensorIntrin& tensor_intrin) {
-  LOG(FATAL) << "TensorIntrin cannot be directly passed to meta schedule. Please register the tensor intrin and pass the intrin name instead.";
+  LOG(FATAL) << "TensorIntrin cannot be directly passed to meta schedule. Please register the "
+                "tensor intrin and pass the intrin name instead.";
 }
 
 /******** Schedule: Annotation ********/