wehub-resource-sync
210 lines
7.0 KiB
Python
210 lines
7.0 KiB
Python
# 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"
|