chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,119 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tile primitive shorthand namespace for TIRx TVMScript."""
+
+import functools
+
+from tvm.tirx import Buffer, BufferRegion
+
+from .builder import tirx as _builder
+
+_TILE_ARG_TYPES = (Buffer, BufferRegion)
+
+
+def _get_arg(args, kwargs, index, name):
+    if len(args) > index:
+        return args[index]
+    return kwargs.get(name)
+
+
+def _require_buffer_arg(op_name, arg_name, value):
+    if not isinstance(value, _TILE_ARG_TYPES):
+        raise TypeError(
+            f"Tx.{op_name} is tile-only and expects `{arg_name}` to be a Buffer "
+            f"or BufferRegion; use T.{op_name} for expression/builtin calls"
+        )
+
+
+def _validate_tile_call(op_name, args, kwargs):
+    dst = _get_arg(args, kwargs, 0, "dst")
+    _require_buffer_arg(op_name, "dst", dst)
+
+    if op_name in {"cast", "max", "min", "permute_layout", "silu"}:
+        src = _get_arg(args, kwargs, 1, "src")
+        _require_buffer_arg(op_name, "src", src)
+    elif op_name in {"sqrt", "exp", "exp2", "reciprocal"}:
+        src = _get_arg(args, kwargs, 1, "src")
+        if src is not None:
+            _require_buffer_arg(op_name, "src", src)
+
+
+def _tile_scoped_op(op_name):
+    scoped_op = getattr(_builder, op_name)
+
+    @functools.wraps(scoped_op._fn)  # pylint: disable=protected-access
+    def wrapper(*args, scope=None, **kwargs):
+        _validate_tile_call(op_name, args, kwargs)
+        return scoped_op._fn(*args, scope=scope, **kwargs)  # pylint: disable=protected-access
+
+    return _builder.ScopedOp(wrapper)
+
+
+_SCOPED_TILE_OP_NAMES = [
+    "add",
+    "binary_chain",
+    "binary_reduce",
+    "cast",
+    "copy",
+    "copy_async",
+    "exp",
+    "exp2",
+    "fdiv",
+    "fill",
+    "fma",
+    "gemm",
+    "gemm_async",
+    "max",
+    "maximum",
+    "memset",
+    "min",
+    "minimum",
+    "mul",
+    "permute_layout",
+    "reciprocal",
+    "reduce_negate",
+    "select",
+    "silu",
+    "sqrt",
+    "sub",
+    "sum",
+    "unary_reduce",
+    "zero",
+]
+
+for _op_name in _SCOPED_TILE_OP_NAMES:
+    globals()[_op_name] = _tile_scoped_op(_op_name)
+
+cluster = _builder.ScopeNamespace("cluster", "cluster")
+cta = _builder.ScopeNamespace("cta", "cta")
+wg = _builder.ScopeNamespace("warpgroup", "wg")
+warpgroup = _builder.ScopeNamespace("warpgroup", "warpgroup")
+warp = _builder.ScopeNamespace("warp", "warp")
+thread = _builder.ScopeNamespace("thread", "thread")
+
+compose_op = _builder.compose_op
+
+__all__ = [
+    *_SCOPED_TILE_OP_NAMES,
+    "cluster",
+    "compose_op",
+    "cta",
+    "thread",
+    "warp",
+    "warpgroup",
+    "wg",
+]