Merge branch 'main' of https://github.com/tile-ai/tilelang into assume_1121

Author: LeiWang1999

src/tl_templates/cuda/atomic.h

@@ -169,6 +169,7 @@ TL_DEVICE T1 AtomicMinRet(T1 &ref, T2 val,
   }
 }
 
+#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 890))
 template <typename T1, typename T2>
 TL_DEVICE void AtomicAdd(T1 &ref, T2 val,
                          int memory_order = int(cuda::memory_order_relaxed)) {
@@ -236,14 +237,18 @@ TL_DEVICE void AtomicAdd(T1 &ref, T2 val,
       }
     }
   } else {
-#if CUDART_VERSION >= 11080
-    cuda::atomic_ref<NT1, cuda::thread_scope_device> aref(*address);
-    aref.fetch_add(cuda_cast<NT1>(val), cuda::memory_order(memory_order));
-#else
-    TL_NOT_IMPLEMENTED();
-#endif
+    atomicAdd(reinterpret_cast<NT1 *>(address), cuda_cast<NT1>(val));
   }
 }
+#else
+template <typename T1, typename T2>
+TL_DEVICE void AtomicAdd(T1 &ref, T2 val,
+                         int memory_order = int(cuda::memory_order_relaxed)) {
+  using NT1 = typename normalize_atomic_type<T1>::type;
+  (void)memory_order;
+  atomicAdd(reinterpret_cast<NT1 *>(&ref), cuda_cast<NT1>(val));
+}
+#endif
 
 template <typename T1, typename T2>
 TL_DEVICE T1 AtomicAddRet(T1 &ref, T2 val,
@@ -643,6 +648,48 @@ AtomicAddx4Ret(float *ref, float *val,
     return ret_val;
   }
 }
+#else
+TL_DEVICE void AtomicAddx2(float *ref, float *val,
+                           int memory_order = int(cuda::memory_order_relaxed)) {
+  (void)memory_order;
+  float2 add_val = *reinterpret_cast<float2 *>(val);
+  atomicAdd(ref + 0, add_val.x);
+  atomicAdd(ref + 1, add_val.y);
+}
+
+TL_DEVICE float2
+AtomicAddx2Ret(float *ref, float *val,
+               int memory_order = int(cuda::memory_order_relaxed)) {
+  (void)memory_order;
+  float2 add_val = *reinterpret_cast<float2 *>(val);
+  float2 ret;
+  ret.x = atomicAdd(ref + 0, add_val.x);
+  ret.y = atomicAdd(ref + 1, add_val.y);
+  return ret;
+}
+
+TL_DEVICE void AtomicAddx4(float *ref, float *val,
+                           int memory_order = int(cuda::memory_order_relaxed)) {
+  (void)memory_order;
+  float4 add_val = *reinterpret_cast<float4 *>(val);
+  atomicAdd(ref + 0, add_val.x);
+  atomicAdd(ref + 1, add_val.y);
+  atomicAdd(ref + 2, add_val.z);
+  atomicAdd(ref + 3, add_val.w);
+}
+
+TL_DEVICE float4
+AtomicAddx4Ret(float *ref, float *val,
+               int memory_order = int(cuda::memory_order_relaxed)) {
+  (void)memory_order;
+  float4 add_val = *reinterpret_cast<float4 *>(val);
+  float4 ret;
+  ret.x = atomicAdd(ref + 0, add_val.x);
+  ret.y = atomicAdd(ref + 1, add_val.y);
+  ret.z = atomicAdd(ref + 2, add_val.z);
+  ret.w = atomicAdd(ref + 3, add_val.w);
+  return ret;
+}
 #endif
 
 template <typename T> TL_DEVICE T AtomicLoad(T &ref, int memory_order) {

testing/python/language/test_tilelang_intimm.py

@@ -0,0 +1,28 @@
+import tilelang
+import tilelang.testing
+import tilelang.language as T
+
+
+def test_tilelang_intimm():
+    T.int32(0x7fffffff)
+    T.int32(-0x7fffffff - 1)
+    T.uint32(0xffffffff)
+    T.int64(0x7fffffffffffffff)
+    T.int64(-0x7fffffffffffffff - 1)
+    T.uint64(0xffffffffffffffff)
+
+    a = T.int32()
+    a & 0x7fffffff
+
+    a = T.uint32()
+    a & 0xffffffff
+
+    a = T.int64()
+    a & 0x7fffffffffffffff
+
+    a = T.uint64()
+    a & T.uint64(0xffffffffffffffff)
+
+
+if __name__ == '__main__':
+    tilelang.testing.main()

testing/python/language/test_tilelang_language_frontend_v2.py

@@ -394,6 +394,64 @@ def prim_call_macro():
     except ValueError:
         pass
 
+    try:
+
+        @T.macro
+        def macro_with_var(x: T.Ref):
+            x = 1  # noqa: F841
+
+        @T.prim_func
+        def prim_call_macro():
+            with T.Kernel(1):
+                x = T.alloc_var(T.int32)
+                macro_with_var(x)
+
+        assert 'x[0] = 1' in prim_call_macro.script()
+    finally:
+        pass
+
+    try:
+
+        @T.macro
+        def macro_with_var(x: T.Ref):
+            x = 1  # noqa: F841
+
+        @T.prim_func
+        def prim_call_macro():
+            with T.Kernel(1):
+                x = 1
+                macro_with_var(x)
+
+        raise RuntimeError("Expect to report an error, x should not be passed as T.Var")
+    except ValueError:
+        pass
+
+
+def frame_inside_macro():
+
+    @tilelang.jit
+    def get_sample_kernel():
+
+        @T.macro
+        def transform(x):
+            return x + 1
+
+        @T.prim_func
+        def sample_kernel(
+            num_blocks: T.int32,
+            idx_out: T.Tensor[(32,), T.int32],
+        ):
+            with T.Kernel(num_blocks, threads=32) as block_idx:  # noqa: F841
+                fragment = T.alloc_fragment(32, 'int32')
+                T.copy(idx_out, fragment)
+
+                for i in T.Parallel(32):
+                    idx_out[i] = transform(fragment[i])
+
+        return sample_kernel
+
+    kernel = get_sample_kernel()  # noqa: F841
+
 
 if __name__ == '__main__':
     tilelang.testing.main()

tilelang/language/__init__.py

@@ -22,6 +22,7 @@
     FragmentBuffer,  # noqa: F401
     SharedBuffer,  # noqa: F401
     LocalBuffer,  # noqa: F401
+    Ref,  # noqa: F401
 )
 from .loop import serial, Parallel, Persistent, Pipelined  # noqa: F401
 from .frame import has_let_value, get_let_value  # noqa: F401

tilelang/language/proxy.py

@@ -1,7 +1,7 @@
 """The language interface for tl programs."""
 from __future__ import annotations
 
-from typing import Any, SupportsIndex, TYPE_CHECKING
+from typing import Any, SupportsIndex, TYPE_CHECKING, Generic, TypeVar
 from collections.abc import Sequence
 from typing_extensions import Self
 
@@ -263,13 +263,21 @@ class SharedBuffer(BaseTensor):
 
     class LocalBuffer(BaseTensor):
         ...
+
+    _T = TypeVar('_T')
+
+    class Ref(Generic[_T], tir.Var):
+        ...
 else:
     Tensor = TensorProxy()  # pylint: disable=invalid-name
     StridedTensor = StridedTensorProxy()  # pylint: disable=invalid-name
     FragmentBuffer = FragmentBufferProxy()  # pylint: disable=invalid-name
     SharedBuffer = SharedBufferProxy()  # pylint: disable=invalid-name
     LocalBuffer = LocalBufferProxy()  # pylint: disable=invalid-name
 
+    class Ref:
+        ...
+
 
 def ptr(dtype: str | None = None,
         storage_scope: str = "global",

tilelang/language/v2/builder.py

@@ -80,6 +80,10 @@ class MacroFrame(Frame):
     ...
 
 
+class ExitedMacroFrame(Frame):
+    ...
+
+
 class BoolOpFrame(Frame):
     ...
 
@@ -164,8 +168,22 @@ def macro(self, name=None, annotations=None):
         save = self.name_inside_frame, self.arg_annotations
         self.name_inside_frame = {}
         self.arg_annotations = annotations or {}
-        with self.with_frame(MacroFrame()):
-            yield
+        pos = len(self.frames)
+        # here we add a ExitedMacroFrame to preserve the frame stack inside macro
+        # because macro may bind some variable, and return it
+        #
+        # ```py
+        # @T.macro
+        # def foo(x):
+        #    y = x + 1
+        #    return y
+        # @T.prim_func
+        # def bar():
+        #    c = foo(1) # macro generates let y = x + 1
+        #    d = c # d = c should lay inside frame of `let y = x + 1`
+        self.frames.append(MacroFrame())
+        yield
+        self.frames[pos] = ExitedMacroFrame()
         self.name_inside_frame, self.arg_annotations = save
 
     def get(self):
@@ -335,7 +353,7 @@ def bind(self, name, value, annot=BaseBuilder.empty):
             assert frame is not None, f"Variable `{name}` is not defined inside any control flow."
             if name in self.name_inside_frame and self.name_inside_frame[name] in self.frames:
                 logger.warning(
-                    f'Variable `{name}` shadows another declared value, Are you forgetting to allocate it as a var?',
+                    f'Variable `{name}` is declared twice, are you looking for a T.alloc_var?',
                     stack_info=True,
                     stacklevel=2,
                 )
@@ -475,7 +493,11 @@ def rval(self, name: str, value: Any) -> Any:
         return self.unwrap_value(value)
 
     def macro_arg(self, name, value):
-        if self.arg_annotations.get(name, None) is Var:
+        from tilelang.language.proxy import Ref
+        annot_value = self.arg_annotations.get(name, None)
+        if annot_value is Var or annot_value is Ref:
+            if annot_value is Var:
+                logger.warning('Use `T.Var` as macro annotations is deprecated, please use `T.Ref`')
             is_var = isinstance(value, tvm.tir.BufferLoad) and value.buffer.scope() == 'local.var'
             if not is_var:
                 raise ValueError(

tilelang/language/v2/dtypes.py

@@ -87,8 +87,12 @@
     'float8_e8m0fnu': 'Float8E8M0FNU'
 }
 
+int_ = int
+
 
 def __dtype_call__(self: dtype, expr=None, is_size_var: bool = False) -> tir.Var:
+    if isinstance(expr, int_):
+        return tvm.tir.const(expr, dtype=self)
     if self in _STR_TO_TVM_DTYPE_CALL:
         attr = _STR_TO_TVM_DTYPE_CALL[self]
         call = getattr(tb_ffi, attr, None)
@@ -151,6 +155,10 @@ class int8(dtype): ...
     class int16(dtype): ...
     class int32(dtype): ...
     class int64(dtype): ...
+    class int8x2(dtype): ...
+    class int16x2(dtype): ...
+    class int32x2(dtype): ...
+    class int64x2(dtype): ...
     class int8x4(dtype): ...
     class int16x4(dtype): ...
     class int32x4(dtype): ...
@@ -175,6 +183,10 @@ class uint8(dtype): ...
     class uint16(dtype): ...
     class uint32(dtype): ...
     class uint64(dtype): ...
+    class uint8x2(dtype): ...
+    class uint16x2(dtype): ...
+    class uint32x2(dtype): ...
+    class uint64x2(dtype): ...
     class uint8x4(dtype): ...
     class uint16x4(dtype): ...
     class uint32x4(dtype): ...
@@ -308,6 +320,10 @@ class bfloat16(dtype): ...
     int16 = dtype('int16')
     int32 = dtype('int32')
     int64 = dtype('int64')
+    int8x2 = dtype('int8x2')
+    int16x2 = dtype('int16x2')
+    int32x2 = dtype('int32x2')
+    int64x2 = dtype('int64x2')
     int8x4 = dtype('int8x4')
     int16x4 = dtype('int16x4')
     int32x4 = dtype('int32x4')
@@ -332,6 +348,10 @@ class bfloat16(dtype): ...
     uint16 = dtype('uint16')
     uint32 = dtype('uint32')
     uint64 = dtype('uint64')
+    uint8x2 = dtype('uint8x2')
+    uint16x2 = dtype('uint16x2')
+    uint32x2 = dtype('uint32x2')
+    uint64x2 = dtype('uint64x2')
     uint8x4 = dtype('uint8x4')
     uint16x4 = dtype('uint16x4')
     uint32x4 = dtype('uint32x4')
@@ -464,6 +484,10 @@ class bfloat16(dtype): ...
     'int16',
     'int32',
     'int64',
+    'int8x2',
+    'int16x2',
+    'int32x2',
+    'int64x2',
     'int8x4',
     'int16x4',
     'int32x4',
@@ -488,6 +512,10 @@ class bfloat16(dtype): ...
     'uint16',
     'uint32',
     'uint64',
+    'uint8x2',
+    'uint16x2',
+    'uint32x2',
+    'uint64x2',
     'uint8x4',
     'uint16x4',
     'uint32x4',