Merge remote-tracking branch 'upstream/main' into update-maint-release

Author: oraluben

.github/workflows/ci.yml

@@ -56,10 +56,10 @@ jobs:
         run: |
           "${{ steps.setup-pylowest.outputs.python-path }}" -m compileall -q -f tilelang
 
-      - name: Setup Python 3.12
+      - name: Setup Python 3.9
         uses: actions/setup-python@v6
         with:
-          python-version: "3.12"
+          python-version: "3.9"
           update-environment: true
           cache: pip
           cache-dependency-path: |

examples/attention_sink/benchmark_gqa_sink_fwd.py

@@ -5,6 +5,7 @@
 import triton.language as tl
 from triton.tools.tensor_descriptor import TensorDescriptor
 from example_gqa_sink_fwd_bhsd_wgmma_pipelined import flashattn, ref_program, gen_inputs
+from typing import Optional
 
 
 @triton.jit
@@ -94,7 +95,7 @@ def triton_kernel(
     Out.store([off_z, off_h, start_m * BLOCK_M, 0], acc)
 
 
-def triton_program(Q, K, V, Sinks, window_size: int | None = None) -> torch.Tensor:
+def triton_program(Q, K, V, Sinks, window_size: Optional[int] = None) -> torch.Tensor:
     bs, n_heads, seq_q, head_dim = Q.shape
     _, n_heads_kv, seq_kv, _ = K.shape
     BLOCK_M = 64
@@ -130,7 +131,7 @@ def main(
     seq_kv: int = 256,
     dim: int = 128,
     groups: int = 8,
-    window_size: int | None = None,
+    window_size: Optional[int] = None,
     dtype: str = "float16",
     tune: bool = False,
 ):

examples/attention_sink/benchmark_mha_sink_fwd.py

@@ -5,6 +5,7 @@
 import triton.language as tl
 from triton.tools.tensor_descriptor import TensorDescriptor
 from example_mha_sink_fwd_bhsd_wgmma_pipelined import flashattn, ref_program, gen_inputs
+from typing import Optional
 
 
 @triton.jit
@@ -93,7 +94,7 @@ def triton_kernel(
     Out.store([off_z, off_h, start_m * BLOCK_M, 0], acc)
 
 
-def triton_program(Q, K, V, Sinks, window_size: int | None = None) -> torch.Tensor:
+def triton_program(Q, K, V, Sinks, window_size: Optional[int] = None) -> torch.Tensor:
     bs, n_heads, seq_q, head_dim = Q.shape
     seq_kv = K.shape[2]
     BLOCK_M = 64
@@ -125,7 +126,7 @@ def main(batch: int = 1,
          seq_q: int = 256,
          seq_kv: int = 256,
          dim: int = 128,
-         window_size: int | None = None,
+         window_size: Optional[int] = None,
          dtype: str = "float16",
          tune: bool = False):
     torch_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]

examples/attention_sink/example_gqa_sink_bwd_bhsd.py

@@ -444,7 +444,7 @@ def main(BATCH: int = 1,
          N_CTX: int = 512,
          D_HEAD: int = 64,
          groups: int = 2,
-         window_size: int | None = None,
+         window_size: Optional[int] = None,
          dtype: str = "float16"):
     torch_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]
     if window_size is not None:

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

@@ -272,7 +272,7 @@ def main(
     seq_kv: int = 256,
     dim: int = 128,
     groups: int = 8,
-    window_size: int | None = None,
+    window_size: Optional[int] = None,
     dtype: str = "float16",
     tune: bool = False,
 ):

examples/attention_sink/example_mha_sink_bwd_bhsd.py

@@ -440,7 +440,7 @@ def main(BATCH: int = 1,
          H: int = 1,
          N_CTX: int = 512,
          D_HEAD: int = 128,
-         window_size: int | None = None,
+         window_size: Optional[int] = None,
          dtype: str = "float16"):
     torch_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]
     if window_size is not None:

examples/attention_sink/example_mha_sink_fwd_bhsd.py

@@ -253,7 +253,7 @@ def main(batch: int = 1,
          seq_q: int = 256,
          seq_kv: int = 256,
          dim: int = 128,
-         window_size: int | None = None,
+         window_size: Optional[int] = None,
          dtype: str = "float16",
          tune: bool = False):
     torch_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py

@@ -263,7 +263,7 @@ def main(batch: int = 1,
          seq_q: int = 256,
          seq_kv: int = 256,
          dim: int = 128,
-         window_size: int | None = None,
+         window_size: Optional[int] = None,
          dtype: str = "float16",
          tune: bool = False):
     torch_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype]

pyproject.toml

@@ -2,7 +2,7 @@
 name = "tilelang"
 description = "A tile level programming language to generate high performance code."
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 authors = [{ name = "TileLang Contributors" }, { name = "Tile-AI" }]
 maintainers = [{ name = "Lei Wang", email = "leiwang1999@outlook.com" }]
 license = "MIT"
@@ -14,7 +14,6 @@ classifiers = [
     "Operating System :: MacOS",
     "Programming Language :: C++",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -128,7 +127,7 @@ skip = [
 ]
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py39"
 line-length = 100
 output-format = "full"
 

src/op/fill.cc

@@ -17,6 +17,7 @@
 #include "../transform/loop_partition.h"
 #include "../transform/loop_vectorize.h"
 #include "builtin.h"
+#include "region.h"
 
 namespace tvm {
 namespace tl {
@@ -62,7 +63,30 @@ using namespace tir;
 Fill::Fill(Array<PrimExpr> args, BufferMap vmap) {
   ObjectPtr<FillNode> node = tvm::ffi::make_object<FillNode>();
 
-  if (args[0]->IsInstance<BufferLoadNode>()) {
+  // Case 1: Region descriptor call (tl.region)
+  if (const auto *call = args[0].as<CallNode>()) {
+    if (call->op.same_as(RegionOp::Get())) {
+      auto region = RegionOp(call->args, vmap);
+      node->dst = region->GetBuffer();
+      node->region = region->GetRanges();
+    } else if (call->op.same_as(builtin::tvm_access_ptr())) {
+      node->dst = vmap[GetVarFromAccessPtr(args[0])];
+      for (int i = 0; i < node->dst->shape.size(); i++) {
+        node->region.push_back(Range(0, node->dst->shape[i]));
+      }
+    } else {
+      ICHECK(false) << "Unsupported call op in tl.fill: "
+                    << Downcast<Op>(call->op)->name;
+    }
+
+    // Case 2: Explicit BufferRegion (legacy path)
+  } else if (args[0]->IsInstance<BufferRegionNode>()) {
+    auto region = Downcast<BufferRegion>(args[0]);
+    node->dst = region->buffer;
+    node->region = region->region;
+
+    // Case 3: Vector/scalar region expressed via BufferLoad indices
+  } else if (args[0]->IsInstance<BufferLoadNode>()) {
     auto buffer_load = Downcast<BufferLoad>(args[0]);
     for (const auto &index : buffer_load->indices) {
       if (const auto *ramp = index.as<RampNode>()) {
@@ -77,6 +101,7 @@ Fill::Fill(Array<PrimExpr> args, BufferMap vmap) {
       }
     }
     node->dst = buffer_load->buffer;
+    // Case 4: Access pointer, fill the full buffer
   } else {
     node->dst = vmap[GetVarFromAccessPtr(args[0])];
     for (int i = 0; i < node->dst->shape.size(); i++) {
@@ -95,14 +120,19 @@ Fill::Fill(Array<PrimExpr> args, BufferMap vmap) {
       << " != " << node->dst->shape.size();
   for (int i = 0; i < node->region.size(); i++) {
     // bound check if region is static
-    if (node->region[i]->min.as<IntImm>()) {
-      int64_t min = Downcast<IntImm>(node->region[i]->min)->value;
+    if (const auto *min_imm = node->region[i]->min.as<IntImmNode>()) {
+      int64_t min = min_imm->value;
       ICHECK_GE(min, 0) << "region[" << i << "] = " << min << " < 0";
     }
-    if (node->region[i]->extent.as<IntImm>()) {
-      int64_t extent = Downcast<IntImm>(node->region[i]->extent)->value;
-      ICHECK_LE(extent, Downcast<IntImm>(node->dst->shape[i])->value)
-          << "region[" << i << "] = " << extent << " > " << node->dst->shape[i];
+    if (const auto *extent_imm = node->region[i]->extent.as<IntImmNode>()) {
+      // Only perform the upper-bound check when the destination shape
+      // extent is also statically known. If the shape is symbolic (e.g., Var),
+      // skip this static check to avoid invalid downcasts.
+      if (const auto *shape_imm = node->dst->shape[i].as<IntImmNode>()) {
+        ICHECK_LE(extent_imm->value, shape_imm->value)
+            << "region[" << i << "] = " << extent_imm->value << " > "
+            << node->dst->shape[i];
+      }
     }
   }
   data_ = std::move(node);
@@ -140,7 +170,8 @@ For FillNode::MakeSIMTLoop(arith::Analyzer *analyzer) const {
   for (int i = 0; i < ndim; i++) {
     Var var = Var(std::string{char('i' + i)}, region[i]->extent->dtype);
     loop_vars.push_back({region[i], var, IterVarType::kDataPar});
-    dst_indices.push_back(var);
+    // Offset the loop induction variable by region min to honor sliced regions
+    dst_indices.push_back(region[i]->min + var);
   }
   Stmt body = BufferStore(dst, value, dst_indices);
   for (int i = ndim - 1; i >= 0; i--) {
@@ -202,6 +233,7 @@ Stmt FillNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
     return vectorized_thread_loop;
   } else {
     LOG(FATAL) << "Unsupported scope " << dst.scope();
+    return Stmt();
   }
 }
 
@@ -229,4 +261,4 @@ TIR_REGISTER_TL_OP(Fill, fill)
 TVM_FFI_STATIC_INIT_BLOCK() { FillNode::RegisterReflection(); }
 
 } // namespace tl
-} // namespace tvm
+} // namespace tvm