[chore] refactor binding

Author: botbw

src/op/gemm_sp_py.cc

@@ -5,6 +5,7 @@
  */
 
 #include "gemm_sp_py.h"
+#include "utils.h"
 
 #include "builtin.h"
 #include <tvm/tir/builtin.h>
@@ -48,17 +49,19 @@ using namespace tir;
  *       fails with an ICHECK (runtime assertion). No other validation is
  *       performed here.
  */
-GemmSPPy::GemmSPPy(Array<PrimExpr> args, BufferMap vmap) {
+GemmSPPy::GemmSPPy(Array<PrimExpr> args) {
   ObjectPtr<GemmSPPyNode> node = tvm::ffi::make_object<GemmSPPyNode>();
 
-  node->Aptr = args[0];
-  node->Eptr = args[1];
-  node->Bptr = args[2];
-  node->Cptr = args[3];
-  node->A = vmap[GetVarFromAccessPtr(node->Aptr)];
-  node->E = vmap[GetVarFromAccessPtr(node->Eptr)];
-  node->B = vmap[GetVarFromAccessPtr(node->Bptr)];
-  node->C = vmap[GetVarFromAccessPtr(node->Cptr)];
+  node->aRegion_ = NormalizeToBufferRegion(args[0]);
+  node->eRegion_ = NormalizeToBufferRegion(args[1]);
+  node->bRegion_ = NormalizeToBufferRegion(args[2]);
+  node->cRegion_ = NormalizeToBufferRegion(args[3]);
+
+  node->A = node->aRegion_->buffer;
+  node->E = node->eRegion_->buffer;
+  node->B = node->bRegion_->buffer;
+  node->C = node->cRegion_->buffer;
+
   node->trans_A = args[4].as<Bool>().value();
   node->trans_B = args[5].as<Bool>().value();
   node->trans_E = args[6].as<Bool>().value();

src/op/gemm_sp_py.h

@@ -23,7 +23,7 @@ class GemmSPPyNode : public TileOperatorNode {
   bool CheckWGMMA() const;
   tir::Buffer A, E, B, C;
   // pointer to the A, E, B, C
-  PrimExpr Aptr, Eptr, Bptr, Cptr;
+  BufferRegion aRegion_, eRegion_, bRegion_, cRegion_;
   bool trans_A, trans_B, trans_E;
   int M, N, K;
   int stride_A, stride_B;
@@ -33,6 +33,8 @@ class GemmSPPyNode : public TileOperatorNode {
   // only will be enabled under cdna mfma instructions
   int kPack = 1;
   int wg_wait = 0;
+
+  // use GemmWarp Policy here as the atom size are flexible in v2
   mutable GemmWarpPolicy policy;
 
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.GemmSPPy", GemmSPPyNode,
@@ -45,10 +47,10 @@ class GemmSPPyNode : public TileOperatorNode {
         .def_ro("E", &GemmSPPyNode::E)
         .def_ro("B", &GemmSPPyNode::B)
         .def_ro("C", &GemmSPPyNode::C)
-        .def_ro("Aptr", &GemmSPPyNode::Aptr)
-        .def_ro("Eptr", &GemmSPPyNode::Eptr)
-        .def_ro("Bptr", &GemmSPPyNode::Bptr)
-        .def_ro("Cptr", &GemmSPPyNode::Cptr)
+        .def_ro("aRegion", &GemmSPPyNode::aRegion_)
+        .def_ro("eRegion", &GemmSPPyNode::eRegion_)
+        .def_ro("bRegion", &GemmSPPyNode::bRegion_)
+        .def_ro("cRegion", &GemmSPPyNode::cRegion_)
         .def_ro("trans_A", &GemmSPPyNode::trans_A)
         .def_ro("trans_B", &GemmSPPyNode::trans_B)
         .def_ro("trans_E", &GemmSPPyNode::trans_E)
@@ -82,7 +84,7 @@ class GemmSPPy : public TileOperator {
 public:
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GemmSPPy, TileOperator,
                                              GemmSPPyNode);
-  TVM_DLL GemmSPPy(Array<PrimExpr> args, BufferMap vmap);
+  TVM_DLL GemmSPPy(Array<PrimExpr> args);
   static const Op &Get();
 };
 

tilelang/language/experimental/gemm_sp.py

@@ -1,10 +1,17 @@
 """The language interface for tl programs."""
 from __future__ import annotations
 from tilelang.primitives.gemm.base import GemmWarpPolicy
-from tilelang.utils.language import get_buffer_region_from_load
 import tilelang.language as T
 from tvm import tir
-from tilelang.utils.language import to_buffer_region
+from tilelang.utils.language import (
+    to_buffer_region,
+    retrieve_shape,
+    retrieve_stride,
+    retrieve_offset,
+    prim_expr_equal,
+)
+from tilelang.language.utils import (
+    buffer_region_to_tile_region,)
 
 
 def gemm_sp(
@@ -144,54 +151,13 @@ def legalize_arguments(arg: tir.Buffer | tir.Var):
     B = legalize_arguments(B)
     C = legalize_arguments(C)
 
-    def retrieve_shape(object: tir.Buffer | tir.BufferRegion) -> list[int]:
-        if isinstance(object, tir.Buffer):
-            return object.shape
-        elif isinstance(object, tir.BufferRegion):
-            region = object.region
-            shape = []
-            for r in region:
-                shape.append(r.extent)
-            return shape
-        elif isinstance(object, tir.BufferLoad):
-            region = get_buffer_region_from_load(object).region
-            shape = []
-            for r in region:
-                shape.append(r.extent)
-            return shape
-        else:
-            raise ValueError(
-                f"Unsupported retrieve_shape argument type: {type(object)} for buffer {object}")
-
-    def retrieve_stride(object: tir.Buffer | tir.BufferRegion) -> list[int]:
-        if isinstance(object, tir.Buffer):
-            strides = []
-            stride = 1
-            for s in reversed(object.shape):
-                strides.insert(0, stride)
-                stride *= s
-            return strides
-        elif isinstance(object, tir.BufferRegion):
-            buffer, _ = object.buffer, object.region
-            strides = []
-            stride = 1
-            for s in reversed(buffer.shape):
-                strides.insert(0, stride)
-                stride *= s
-            return strides
-        elif isinstance(object, tir.BufferLoad):
-            buffer = object.buffer
-            strides = []
-            stride = 1
-            for s in reversed(buffer.shape):
-                strides.insert(0, stride)
-                stride *= s
-            return strides
-        else:
-            raise ValueError(
-                f"Unsupported retrieve_stride argument type: {type(object)} for buffer {object}")
+    A_region = to_buffer_region(A_sparse)
+    E_region = to_buffer_region(E)
+    B_region = to_buffer_region(B)
+    C_region = to_buffer_region(C)
 
     A_shape = retrieve_shape(A_sparse)
+    E_shape = retrieve_shape(E)  # nolint: F841
     B_shape = retrieve_shape(B)
     C_shape = retrieve_shape(C)
 
@@ -213,86 +179,30 @@ def retrieve_stride(object: tir.Buffer | tir.BufferRegion) -> list[int]:
     M, N = C_shape
     K = 2 * (A_shape[-2] if transpose_A else A_shape[-1])
     K_B = B_shape[-1] if transpose_B else B_shape[-2]
-    assert K == K_B, f"T.gemm_sp K shape check failed: K_A (wo sparse) = {K}, K_B = {K_B}"
+    assert prim_expr_equal(
+        K, K_B), f"T.gemm_sp K shape check failed: K_A (wo sparse) = {K}, K_B = {K_B}"
 
     stride_a = A_stride[-2]
     stride_b = B_stride[-2]
 
-    def retrieve_ptr(object: tir.Buffer | tir.BufferRegion, access_type: str = "r") -> tir.PrimExpr:
-        if isinstance(object, tir.Buffer):
-            return object.access_ptr(access_type)
-        elif isinstance(object, tir.BufferRegion):
-            buffer, region = object.buffer, object.region
-            indices = []
-            for r in region:
-                indices.append(r.min)
-            strides = []
-            stride = 1
-            for s in reversed(buffer.shape):
-                strides.insert(0, stride)
-                stride *= s
-            offset = 0
-            # not offset the last two dimension
-            for i in range(len(indices) - 2):
-                offset += indices[i] * strides[i]
-            return buffer.access_ptr(access_mask=access_type, offset=offset)
-        elif isinstance(object, tir.BufferLoad):
-            buffer = object.buffer
-            region = get_buffer_region_from_load(object).region
-            indices = []
-            for r in region:
-                indices.append(r.min)
-            strides = []
-            stride = 1
-            for s in reversed(buffer.shape):
-                strides.insert(0, stride)
-                stride *= s
-            offset = 0
-            for i in range(len(indices) - 2):
-                offset += indices[i] * strides[i]
-            return buffer.access_ptr(access_mask=access_type, offset=offset)
-        else:
-            raise ValueError(
-                f"Unsupported retrieve_ptr argument type: {type(object)} for buffer {object}")
-
-    def retrieve_offset(object: tir.Buffer | tir.BufferRegion) -> tir.PrimExpr:
-        """Retrieve the offset of the buffer or buffer region."""
-        if isinstance(object, tir.Buffer):
-            return [0] * len(object.shape)
-        elif isinstance(object, tir.BufferRegion):
-            _, region = object.buffer, object.region
-            indices = []
-            for r in region:
-                indices.append(r.min)
-            return indices
-        elif isinstance(object, tir.BufferLoad):
-            region = get_buffer_region_from_load(object).region
-            indices = []
-            for r in region:
-                indices.append(r.min)
-            return indices
-        else:
-            raise ValueError(
-                f"Unsupported retrieve_offset argument type: {type(object)} for buffer {object}")
-
     A_offset = retrieve_offset(A_sparse)
     B_offset = retrieve_offset(B)
     assert A_offset[-2] == 0, "The offset of the first dimension of A must be 0"
     assert B_offset[-2] == 0, "The offset of the first dimension of B must be 0"
     offset_a = A_offset[-1]
     offset_b = B_offset[-1]
 
-    Aptr = retrieve_ptr(A_sparse, "r")
-    Eptr = retrieve_ptr(E, "r")
-    Bptr = retrieve_ptr(B, "r")
-    Cptr = retrieve_ptr(C, "rw")
+    A_arg = buffer_region_to_tile_region(A_region, "r", [r for r in A_shape])
+    E_arg = buffer_region_to_tile_region(E_region, "r", [r for r in E_shape])
+    B_arg = buffer_region_to_tile_region(B_region, "r", [r for r in B_shape])
+    C_arg = buffer_region_to_tile_region(C_region, "rw", [r for r in C_shape])
     return tir.call_intrin(
         "handle",
         tir.op.Op.get("tl.gemm_sp_py"),
-        Aptr,
-        Eptr,
-        Bptr,
-        Cptr,
+        A_arg,
+        E_arg,
+        B_arg,
+        C_arg,
         transpose_A,
         transpose_B,
         transpose_E,

tilelang/tileop/gemm_sp/gemm_sp_base.py

@@ -83,20 +83,20 @@ def C(self) -> tir.Buffer:
         return self.gemm_sp_node.C
 
     @property
-    def APtr(self) -> tir.PrimExpr:
-        return self.gemm_sp_node.APtr
+    def ARegion(self) -> tir.PrimExpr:
+        return self.gemm_sp_node.ARegion
 
     @property
-    def EPtr(self) -> tir.PrimExpr:
-        return self.gemm_sp_node.EPtr
+    def ERegion(self) -> tir.PrimExpr:
+        return self.gemm_sp_node.ERegion
 
     @property
-    def BPtr(self) -> tir.PrimExpr:
-        return self.gemm_sp_node.BPtr
+    def BRegion(self) -> tir.PrimExpr:
+        return self.gemm_sp_node.BRegion
 
     @property
-    def CPtr(self) -> tir.PrimExpr:
-        return self.gemm_sp_node.CPtr
+    def CRegion(self) -> tir.PrimExpr:
+        return self.gemm_sp_node.CRegion
 
     @property
     def stride_A(self) -> int: