[Refactor] Support python reflection for tile operators (#783)

* Implement Fill operator and related reflection methods in TileLang

- Added Fill operator implementation in `fill.cc` and `fill.h` for element-wise filling of buffers.
- Introduced reflection methods for Fill, AtomicAdd, Copy, Conv2DIm2Col, FinalizeReducer, Gemm, and Parallel operators to enhance introspection capabilities.
- Updated relevant files to register reflection methods and ensure proper initialization in static blocks.
- Removed outdated comments and unnecessary code in various operator files to improve clarity and maintainability.
- Added new Python bindings for the Fill operator in `tilelang/ir/fill.py` and updated the module imports accordingly.

* Refactor operator reflection methods and improve code clarity

- Updated reflection methods for AtomicAdd, Copy, FinalizeReducer, Gemm, and Parallel operators to enhance readability by using `empty()` instead of size checks.
- Consolidated static initialization blocks for various operators to a single line for improved consistency.
- Cleaned up whitespace and formatting in multiple files to adhere to coding standards and improve maintainability.
- Added new Python bindings for operators in the `tilelang/ir` module, ensuring proper registration and organization of imports.

* Refactor GEMM and AtomicAdd operations for improved clarity

- Updated the `GetArchInt` function in `atomic_add.cc` to use `std::string` and `std::stoi` for better readability and type safety.
- Removed unnecessary variables and comments in `gemm_sp.cc` and `gemm.cc` to streamline the `ComputeWarpPartition` method.
- Cleaned up the `layout_reducer.cc` file by removing unused variable declarations, enhancing code clarity.
- Added import for the `ir` module in `tilelang/__init__.py` to ensure proper organization of module imports.

* Remove deprecated operator files from the tilelang IR module

- Deleted files for Fill, AtomicAdd, Copy, Gemm, GemmSP, FinalizeReducer, Parallel, Reduce, and Region operators to streamline the codebase.
- This cleanup enhances maintainability by removing unused code and improving overall organization of the module.

* Refactor imports in tilelang IR module for improved organization

- Updated import statements in `tilelang/ir.py` to reflect changes in the TVM library structure, enhancing clarity and maintainability of the codebase.

* lint fix

* Refactor GEMM and GEMM-SP operations to enhance clarity and maintainability

- Updated the `Gemm` and `GemmSP` classes to utilize a new `GemmWarpPolicy` object for warp partitioning, improving encapsulation and readability.
- Removed deprecated `ComputeWarpPartition` methods and replaced them with calls to the new policy object, streamlining the code.
- Cleaned up comments and unnecessary code in `gemm.cc`, `gemm_sp.cc`, and related header files to enhance overall clarity.
- Introduced a new `GemmWarpPolicyNode` class to manage warp policy attributes and methods, facilitating better organization of related functionalities.
- Updated reflection methods to include the new policy structure, ensuring proper registration and introspection capabilities.

* Refactor Reduce operation to utilize ReduceType class for improved clarity and maintainability

- Replaced multiple conditional checks for reduce types with a single ReduceType object, simplifying the code structure.
- Introduced a new ReduceTypeNode class to encapsulate reduce type logic and methods, enhancing organization.
- Updated MakeInitValue, MakeReduce, and Lower methods to leverage the new ReduceType class, improving readability.
- Added Python bindings for the ReduceType class in tilelang IR module to ensure proper registration and usability.

* comment

* Refactor operator header files for improved readability

- Cleaned up formatting and whitespace in `atomic_add.h`, `copy.h`, `fill.h`, `reduce.cc`, and `reduce.h` to enhance code clarity.
- Consolidated comments and adjusted line breaks for better organization and maintainability across multiple operator definitions.

* Refactor MakeReduce method in ReduceOpNode for clarity

- Updated the parameter name in the MakeReduce method from `rhs` to `b` and assigned it to `rhs` for improved readability.
- This change enhances the clarity of the method's purpose and aligns with the overall refactoring efforts in the Reduce operation.

* Update Reduce operation type checks for consistency

- Changed string comparisons for reduce types in the MakeReduce method from "abs_sum" to "abssum" and "abs_max" to "absmax" for uniformity.
- This adjustment enhances the clarity and consistency of the reduce type handling in the codebase.

Author: LeiWang1999

src/op/atomic_add.cc

@@ -37,9 +37,9 @@ static int GetArchInt(Target target) {
   int arch_int = 0;
   auto s = target->GetAttr<String>("arch");
   ICHECK(s.defined());
-  const char *arch_str = s.value().c_str();
-  if (arch_str[0] == 's' && arch_str[1] == 'm' && arch_str[2] == '_') {
-    arch_int = atoi(&arch_str[3]);
+  std::string arch = s.value();
+  if (arch.rfind("sm_", 0) == 0) {
+    arch_int = std::stoi(arch.substr(3));
   } else {
     arch_int = 0;
   }
@@ -255,7 +255,7 @@ PrimExpr AtomicAddNode::MakePredicate(arith::Analyzer *analyzer,
  */
 For AtomicAddNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   Array<IterVar> loop_vars = MakeIterVars();
-  bool is_scalar = loop_vars.size() == 0;
+  bool is_scalar = loop_vars.empty();
   if (is_scalar) {
     return For(Var("i"), 0, 1, ForKind::kSerial,
                BufferStore(dst, BufferLoad(src, {0}), {0}));
@@ -425,5 +425,7 @@ TIR_REGISTER_TL_OP(AtomicAdd, atomicadd)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));
 
+TVM_FFI_STATIC_INIT_BLOCK({ AtomicAddNode::RegisterReflection(); });
+
 } // namespace tl
 } // namespace tvm

src/op/atomic_add.h

@@ -1,7 +1,6 @@
 /*!
  * \file tl/op/atomic_add.h
- * \brief Define atomic add operator.
- *
+ * \brief Atomic addition operations for concurrent memory updates
  */
 
 #ifndef TVM_TL_OP_ATOMIC_ADD_H_
@@ -10,91 +9,20 @@
 #include "operator.h"
 #include "parallel.h"
 
-/**
- * Lower this tile operator into a TIR statement for the given lowering context.
- *
- * @param T Lowering context containing mapped buffers and iteration
- * information.
- * @param analyzer Arithmetic analyzer used to simplify and reason about
- * expressions.
- * @return A TIR Stmt that implements the atomic-add tile operation for the
- * provided context.
- */
-/**
- * Infer memory/layout mapping for tensors and buffers used by this operator.
- *
- * @param T Layout inference context providing buffer and shape information.
- * @param level Inference aggressiveness level; higher levels may perform more
- * speculative decisions.
- * @return A LayoutMap describing inferred layouts for the operator's inputs and
- * outputs.
- */
-/**
- * Get the Op registration that identifies this tile operator.
- *
- * @return A reference to the registered Op representing this operator.
- */
-/**
- * Create a deep copy of this tile operator node wrapped as a TileOperator.
- *
- * @return A TileOperator handle owning a cloned AtomicAddNode.
- */
-/**
- * Construct a SIMT-style For loop nest (thread/block mapping) appropriate for
- * the operator.
- *
- * @param analyzer Arithmetic analyzer used to simplify loop bounds and
- * predicates.
- * @return A For loop node representing the SIMT-parallel loop structure.
- */
-/**
- * Create iteration variables used by this operator's loop nest.
- *
- * @return An array of IterVar objects describing the loop iteration axes.
- */
-/**
- * Produce index expressions for either source or destination buffer access
- * based on iteration vars.
- *
- * @param ivs IterVars created by MakeIterVars().
- * @param src_dst Selects which indices to produce: 0 for source indices, 1 for
- * destination indices.
- * @return An array of PrimExpr index expressions suitable for indexing the
- * selected buffer.
- */
-/**
- * Build a predicate expression that guards out-of-bounds or conditional
- * accesses for src or dst.
- *
- * @param analyzer Arithmetic analyzer used to simplify the predicate.
- * @param ivs IterVars created by MakeIterVars().
- * @param extents The loop extents corresponding to the itervars.
- * @param src_dst Selects which side the predicate is for: 0 for source, 1 for
- * destination.
- * @return A PrimExpr boolean predicate that evaluates to true for valid
- * iterations.
- */
-/**
- * Construct an AtomicAdd tile operator from operation arguments and a buffer
- * mapping.
- *
- * @param args Operation arguments (e.g., values or indices) specific to the
- * atomic-add semantics.
- * @param vmap Mapping from buffer names to Buffer objects used by this
- * operator.
- */
 namespace tvm {
 namespace tl {
 
 using namespace tir;
 
+/// Node class for atomic addition operations
 class AtomicAddNode : public TileOperatorNode {
 public:
-  Buffer src, dst;
-  Array<Range> src_range, dst_range;
-  IntImm coalesced_width;
+  Buffer src, dst; ///< Source and destination buffers
+  Array<Range> src_range,
+      dst_range;          ///< Access ranges for source and destination
+  IntImm coalesced_width; ///< Width for memory coalescing optimization
 
-  mutable ParallelOp par_op_;
+  mutable ParallelOp par_op_; ///< Associated parallel operation
   static constexpr const char *_type_key = "tl.AtomicAdd";
   TVM_DECLARE_FINAL_OBJECT_INFO(AtomicAddNode, TileOperatorNode);
 
@@ -104,18 +32,47 @@ class AtomicAddNode : public TileOperatorNode {
   static const Op &Get();
   TileOperator Clone() const;
 
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<AtomicAddNode>()
+        .def_ro("src", &AtomicAddNode::src)
+        .def_ro("dst", &AtomicAddNode::dst)
+        .def_ro("src_range", &AtomicAddNode::src_range)
+        .def_ro("dst_range", &AtomicAddNode::dst_range)
+        .def_ro("coalesced_width", &AtomicAddNode::coalesced_width);
+  }
+
+  bool SEqualReduce(const AtomicAddNode *other, SEqualReducer equal) const {
+    return equal(src, other->src) && equal(dst, other->dst) &&
+           equal(src_range, other->src_range) &&
+           equal(dst_range, other->dst_range) &&
+           equal(coalesced_width, other->coalesced_width);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(src);
+    hash_reduce(dst);
+    hash_reduce(src_range);
+    hash_reduce(dst_range);
+    hash_reduce(coalesced_width);
+  }
+
+  static constexpr bool _type_has_method_sequal_reduce = true;
+  static constexpr bool _type_has_method_shash_reduce = true;
+
 protected:
+  /// Create SIMT-style parallel loop structure
   For MakeSIMTLoop(arith::Analyzer *analyzer) const;
+  /// Generate iteration variables for loop nest
   Array<IterVar> MakeIterVars() const;
-
-  // ivs: itervars returned by MakeIterVars()
-  // src_dst: 0 for src_indices, 1 for dst_indices
+  /// Generate buffer indices from iteration variables
   Array<PrimExpr> MakeIndices(const Array<IterVar> &ivs, int src_dst) const;
-
+  /// Create boundary predicate for memory safety
   PrimExpr MakePredicate(arith::Analyzer *analyzer, const Array<IterVar> &ivs,
                          Array<PrimExpr> extents, int src_dst) const;
 };
 
+/// Wrapper class for atomic addition operations
 class AtomicAdd : public TileOperator {
 public:
   TVM_DEFINE_OBJECT_REF_METHODS(AtomicAdd, TileOperator, AtomicAddNode);

src/op/copy.cc

@@ -297,7 +297,7 @@ PrimExpr CopyNode::MakePredicate(arith::Analyzer *analyzer,
  */
 For CopyNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   Array<IterVar> loop_vars = MakeIterVars();
-  bool is_scalar = loop_vars.size() == 0;
+  bool is_scalar = loop_vars.empty();
   if (is_scalar) {
     return For(Var("i"), 0, 1, ForKind::kSerial,
                BufferStore(dst, BufferLoad(src, {0}), {0}));
@@ -1197,7 +1197,7 @@ Stmt CopyNode::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer,
     int swizzle;
     int max_dim;
   };
-  static const SwizzleCheck swizzle_checks[] = {
+  static const std::vector<SwizzleCheck> swizzle_checks = {
       {static_cast<int>(CU_TENSOR_MAP_SWIZZLE_32B), 32},
       {static_cast<int>(CU_TENSOR_MAP_SWIZZLE_64B), 64},
       {static_cast<int>(CU_TENSOR_MAP_SWIZZLE_128B), 128},
@@ -1559,5 +1559,9 @@ TIR_REGISTER_TL_OP(Conv2DIm2ColOp, c2d_im2col)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));
 
+TVM_FFI_STATIC_INIT_BLOCK({
+  CopyNode::RegisterReflection();
+  Conv2DIm2ColOpNode::RegisterReflection();
+});
 } // namespace tl
 } // namespace tvm