chore: import upstream snapshot with attribution

python/tvm/tirx/operator/tile_primitive/dispatch_context.py

@@ -0,0 +1,209 @@
+# 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.
+"""TIRx operator dispatch context."""
+
+from tvm_ffi import register_object
+
+from tvm.ir import Range
+from tvm.runtime import Object, Scriptable
+from tvm.target import Target
+from tvm.tirx import Buffer, IterVar, Stmt, Var, _ffi_api
+from tvm.tirx.exec_scope import ExecScope
+
+
+@register_object("tirx.DispatchContext")
+class DispatchContext(Object, Scriptable):
+    """DispatchContext node.
+
+    Parameters
+    ----------
+    target : Target
+        The target of the dispatch context.
+
+    exec_scope : ExecScope
+        The execution scope of the dispatch context.
+
+    launch_params : Dict[str, Expr]
+        The launch parameters of the dispatch context.
+
+    var_range_map : Dict[Var, Range]
+        A map from loop variables to their ranges.
+
+    callbacks : Dict[str, Object]
+        The callbacks of the dispatch context.
+
+    shared_state : Dict[str, Object]
+        Shared state persisting across dispatch calls within a single lowering pass.
+    """
+
+    target: Target
+    exec_scope: ExecScope
+    launch_params: dict[str, IterVar]
+    var_range_map: dict[Var, Range]
+    alloc_only: bool
+    callbacks: dict[str, Object]
+    shared_state: dict[str, Object]
+    inter: dict[str, list]
+    intra: dict[str, list]
+    scope_kind: str
+
+    kPrivateAlloc = "private_alloc"
+    kDeviceInitStmt = "device_init_stmt"
+    kHostInitStmt = "host_init_stmt"
+    kPostBufferDefStmt = "post_buffer_def_stmt"
+
+    def __init__(
+        self,
+        target: Target,
+        exec_scope: ExecScope,
+        launch_params: dict[str, IterVar],
+        var_range_map: dict[Var, Range],
+        alloc_only: bool = False,
+        callbacks: dict[str, Object] = {},
+        shared_state: dict[str, Object] = {},
+        inter: dict[str, list] | None = None,
+        intra: dict[str, list] | None = None,
+        scope_kind: str = "",
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DispatchContext,  # pylint: disable=no-member
+            target,
+            exec_scope,
+            launch_params,
+            var_range_map,
+            alloc_only,
+            callbacks,
+            shared_state,
+            inter or {},
+            intra or {},
+            scope_kind,
+        )
+
+    def add_alloc_buffer(self, buffer: Buffer) -> None:
+        """Add an allocated buffer to the dispatch context.
+           Can be called only if alloc_only is True.
+           The buffer will be added to the workspace of operator (the key in the workspace is the buffer name).
+
+        Parameters
+        ----------
+        buffer : Buffer
+            The buffer to be added.
+        """  # noqa: E501
+        _ffi_api.DispatchContextAddAllocBuffer(self, buffer)  # pylint: disable=no-member
+
+    def add_init_stmt(self, stmt: Stmt, host: bool = False) -> None:
+        """Add an initialization statement to the dispatch context.
+           Device initialization statements is only allowed if alloc_only is True.
+           Host initialization statements will be ignored if alloc_only is True.
+           The statements will be added to the beginning of the kernel.
+
+        Parameters
+        ----------
+        stmt : Stmt
+            The initialization statement to be added.
+        host : bool
+            Whether the statement is a host statement.
+            If True, the statement will be added to the host code (before the kernel).
+            If False, the statement will be added to the kernel body (at the beginning of the kernel).
+        """  # noqa: E501
+        _ffi_api.DispatchContextAddInitStmt(self, stmt, host)  # pylint: disable=no-member
+
+    def add_post_buffer_def_stmt(self, buffer: Buffer, stmt: Stmt) -> None:
+        """Add a statement to be inserted after a buffer's definition (DeclBuffer/AllocBuffer).
+
+        Parameters
+        ----------
+        buffer : Buffer
+            The buffer whose definition scope the statement should appear in.
+        stmt : Stmt
+            The statement to be inserted.
+        """
+        _ffi_api.DispatchContextAddPostBufferDefStmt(self, buffer, stmt)  # pylint: disable=no-member
+
+    def cache_get(self, key: str) -> Object | None:
+        """Look up a cached value by key.
+
+        Parameters
+        ----------
+        key : str
+            Cache key (built by the caller from construction parameters).
+
+        Returns
+        -------
+        Optional[Object]
+            The cached value, or None on miss.
+        """
+        return _ffi_api.DispatchContextSharedStateGet(self, key)
+
+    def cache_set(self, key: str, value: Object) -> None:
+        """Store a value in the cross-dispatch cache.
+
+        Parameters
+        ----------
+        key : str
+            Cache key (built by the caller from construction parameters).
+        value : Object
+            The object to cache (e.g. a Buffer or Var).
+        """
+        _ffi_api.DispatchContextSharedStateSet(self, key, value)
+
+    def is_cuda(self) -> bool:
+        """Check if the target is CUDA."""
+        return self.target.kind.name == "cuda"
+
+    def is_trn(self) -> bool:
+        """Check if the target is Trainium."""
+        return self.target.kind.name == "trn"
+
+    def is_target(self, name: str) -> bool:
+        """Check if the target kind matches ``name``."""
+        return self.target.kind.name == name
+
+    # -- scope predicates ----------------------------------------------------
+    #
+    # Each ``is_<scope>`` returns True iff the op site is at that scope kind.
+    # Backed by ``self.scope_kind``, which 1-1 maps to a canonical intra
+    # TileLayout shape:
+    #   thread     -> {}
+    #   warp       -> {laneid}
+    #   warpgroup  -> {laneid, wid_in_wg}
+    #   cta        -> {laneid, warpid}
+    #   cluster    -> {laneid, warpid, cta_id}
+    #
+    # Prefer these predicates over raw ``self.scope_kind == "..."`` comparisons
+    # so dispatchers that later need stricter intra/inter shape checks can
+    # tighten the predicate body without touching every call site.
+
+    @property
+    def is_thread(self) -> bool:
+        return self.scope_kind == "thread"
+
+    @property
+    def is_warp(self) -> bool:
+        return self.scope_kind == "warp"
+
+    @property
+    def is_warpgroup(self) -> bool:
+        return self.scope_kind == "warpgroup"
+
+    @property
+    def is_cta(self) -> bool:
+        return self.scope_kind == "cta"
+
+    @property
+    def is_cluster(self) -> bool:
+        return self.scope_kind == "cluster"