remove rubbish

Author: SiriusNEO

src/op/reduce.cc

@@ -82,27 +82,34 @@ static PrimExpr MakeAccessPtrFromRegion(const BufferRegion &region,
                                         int rw_mask) {
   Buffer buf = region->buffer;
   int ndim = static_cast<int>(buf->shape.size());
-  ICHECK(ndim >= 2) << "GEMM expects buffers with at least 2 dims";
-
-  // Compute row-major strides
-  std::vector<PrimExpr> strides(ndim);
-  PrimExpr one = make_const(buf->shape[0].dtype(), 1);
-  PrimExpr cur = one;
-  for (int i = ndim - 1; i >= 0; --i) {
-    strides[i] = cur;
-    cur = cur * buf->shape[i];
-  }
+  ICHECK(ndim == 1 || ndim == 2) << "Cumsum expects buffers with 1 or 2 dims";
+
+  PrimExpr offset, extent;
+  if (ndim == 1) {
+    // Simple 1D region: offset and extent come from the single axis.
+    auto axis = region->region[0];
+    offset = axis->min;
+    extent = axis->extent;
+  } else {
+    // Compute row-major strides for ndim >= 2
+    std::vector<PrimExpr> strides(ndim);
+    PrimExpr one = make_const(buf->shape[0].dtype(), 1);
+    PrimExpr cur = one;
+    for (int i = ndim - 1; i >= 0; --i) {
+      strides[i] = cur;
+      cur = cur * buf->shape[i];
+    }
+    // Offset: sum_{i in [0..ndim-3]} min_i * stride_i
+    offset = make_const(buf->shape[0].dtype(), 0);
+    for (int i = 0; i < ndim - 2; ++i) {
+      offset = offset + region->region[i]->min * strides[i];
+    }
 
-  // Offset: sum_{i in [0..ndim-3]} min_i * stride_i
-  PrimExpr offset = make_const(buf->shape[0].dtype(), 0);
-  for (int i = 0; i < ndim - 2; ++i) {
-    offset = offset + region->region[i]->min * strides[i];
+    // Extent: last two extents product (elements)
+    extent =
+        region->region[ndim - 2]->extent * region->region[ndim - 1]->extent;
   }
 
-  // Extent: last two extents product (elements)
-  PrimExpr extent =
-      region->region[ndim - 2]->extent * region->region[ndim - 1]->extent;
-
   // ptype and return handle
   PrimExpr ptype = tir::TypeAnnotation(buf->dtype);
   Array<PrimExpr> acc_args{ptype, buf->data, offset, extent,

testing/python/language/test_tilelang_language_atomic_add.py

@@ -260,7 +260,7 @@ def test_atomic_addx2():
     run_atomic_addx2(32, 64, 8, 16)
 
 
-@tilelang.jit(debug_root_path="./testing/python/language")
+@tilelang.jit()
 def atomic_different_memory_orders_program(M, N, block_M, block_N, dtype="float"):
 
     @T.prim_func

tilelang/analysis/ast_printer.py

@@ -1,6 +1,5 @@
 from tvm import tir
-from tvm.tir import (
-    PrimFunc,)
+from tvm.tir import PrimFunc
 from tvm.tir.transform import prim_func_pass
 from tvm.tir.stmt_functor import ir_transform
 

tilelang/language/reduce.py

@@ -246,8 +246,8 @@ def cumsum_fragment(src: tir.Buffer, dst: tir.Buffer, dim: int, reverse: bool) -
     tir.call_intrin(
         "handle",
         tir.op.Op.get("tl.cumsum"),
-        buffer_to_tile_region(cumsum_smem),
-        buffer_to_tile_region(cumsum_smem),
+        buffer_to_tile_region(cumsum_smem, "r"),
+        buffer_to_tile_region(cumsum_smem, "w"),
         dim,
         reverse,
     )