gemini-code-assist[bot] commented on code in PR #19757: URL: https://github.com/apache/tvm/pull/19757#discussion_r3408320078
########## python/tvm/backend/__init__.py: ########## @@ -0,0 +1,203 @@ +# 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. +"""Backend-owned Python modules and load hooks.""" + +from __future__ import annotations + +import importlib +import importlib.util +import sys +import types +from pkgutil import extend_path +from typing import Any + +__path__ = extend_path(__path__, __name__) # type: ignore[name-defined] + +_BUILTIN_BACKENDS = ( + "cuda", + "metal", + "rocm", + "trn", + "opencl", + "vulkan", + "webgpu", + "hexagon", + "adreno", +) +_LOADED_BACKENDS: dict[str, Any] = {} + + +class _AliasModule(types.ModuleType): + """Module object that exposes a backend module under a public alias.""" + + def __init__(self, fullname: str, module): + super().__init__(fullname, getattr(module, "__doc__", None)) + self.__dict__["__tvm_backend_module__"] = module + self.__dict__["__package__"] = fullname.rpartition(".")[0] + if hasattr(module, "__all__"): + self.__dict__["__all__"] = module.__all__ + if hasattr(module, "__path__"): + self.__dict__["__path__"] = [] + + def __getattr__(self, name: str): + return getattr(self.__dict__["__tvm_backend_module__"], name) + + def __dir__(self): + return sorted(set(super().__dir__()) | set(dir(self.__dict__["__tvm_backend_module__"]))) Review Comment:  `_AliasModule` does not override `__setattr__` and `__delattr__`, which causes attributes set on the aliased module to not be reflected in the underlying module. Please override them to forward writes and deletes to the underlying module. ```python def __getattr__(self, name: str): return getattr(self.__dict__["__tvm_backend_module__"], name) def __setattr__(self, name: str, value: Any): if name in ("__tvm_backend_module__", "__package__", "__all__", "__path__"): super().__setattr__(name, value) else: setattr(self.__dict__["__tvm_backend_module__"], name, value) def __delattr__(self, name: str): if name in ("__tvm_backend_module__", "__package__", "__all__", "__path__"): super().__delattr__(name) else: delattr(self.__dict__["__tvm_backend_module__"], name) def __dir__(self): return sorted(set(super().__dir__()) | set(dir(self.__dict__["__tvm_backend_module__"])))" ``` ########## python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py: ########## @@ -516,80 +521,152 @@ def compute_canonical_params(buf, buf_region, dtype, is_transposed): Tuple of (swizzle_mode, ldo, sdo, is_mn_major). """ region = list(buf_region.region) Review Comment:  The `analyzer` variable is used on line 614 but is not defined within the scope of `compute_canonical_params`, which will raise a `NameError` at runtime. Please instantiate `analyzer = Analyzer()` at the beginning of the function. ```suggestion region = list(buf_region.region) analyzer = Analyzer() ``` ########## python/tvm/backend/cuda/lang/alloc_pool.py: ########## @@ -0,0 +1,529 @@ +# 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. +"""SMEM and TMEM bump-allocator pools for TIRX kernels.""" + +from __future__ import annotations + +import functools +import operator + +from tvm import DataType +from tvm.tirx.layout import S, TCol, TileLayout, TLane + +# --------------------------------------------------------------------------- +# ir_builder helpers — imported lazily to avoid circular deps at module level +# --------------------------------------------------------------------------- + +_ir = None + + +def _get_ir(): + global _ir + if _ir is None: + from tvm.tirx.script.builder import ir as _mod + + _ir = _mod + return _ir + + +def _get_frame(): + from tvm.tirx.script.builder import frame + + return frame + + +# --------------------------------------------------------------------------- +# Shared utilities +# --------------------------------------------------------------------------- + +_POOL_UNSET = object() + + +def _default_tmem_layout(rows, cols): + return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)]) + + +def _emit_stmt(expr): + ir = _get_ir() + ir.add_to_parent(ir.evaluate(expr)) + + +def _shape_product(shape): + return functools.reduce(operator.mul, shape, 1) + + +def _auto_swizzle_mode(dtype): + """Select the default MMA swizzle mode for a shared-memory allocation.""" + from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode + + del dtype + return SwizzleMode.SWIZZLE_128B_ATOM + + +def _swizzle_atom_bytes(swizzle_mode): + """Return the row width (in bytes) of one swizzle atom for *swizzle_mode*.""" + from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode + + return { + SwizzleMode.SWIZZLE_NONE: 0, + SwizzleMode.SWIZZLE_32B_ATOM: 32, + SwizzleMode.SWIZZLE_64B_ATOM: 64, + SwizzleMode.SWIZZLE_128B_ATOM: 128, + }[swizzle_mode] + + +def _suggest_swizzle_for_row_bytes(row_bytes): + """Pick the largest valid swizzle mode whose atom row fits within *row_bytes*.""" + + for atom_bytes, mode in ( + (128, "SWIZZLE_128B_ATOM"), + (64, "SWIZZLE_64B_ATOM"), + (32, "SWIZZLE_32B_ATOM"), + ): + if row_bytes >= atom_bytes and row_bytes % atom_bytes == 0: + return mode + return "SWIZZLE_NONE" + + +def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): + """Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible. + + ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]`` + over the last two logical dimensions of *shape*. If the row width or row count of + the request is smaller than (or not a multiple of) the atom, the underlying + ``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an + opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the + misconfiguration here so callers see *what* is wrong and *how* to fix it. + + Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom). + """ + from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode + + if swizzle_mode == SwizzleMode.SWIZZLE_NONE: + return + + if len(shape) < 2: + raise ValueError( + f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; " + f"swizzled MMA layouts tile over the last two dims (rows, cols). " + f"Use swizzle_mode='none' for 1-D allocations." + ) + + # Only validate concrete int dims; symbolic dims fall through (the analyzer + # in C++ will still ICHECK on them, but at least we don't false-positive). + rows = shape[-2] + cols = shape[-1] + if not (isinstance(rows, int) and isinstance(cols, int)): + return + + dtype_bytes = DataType(dtype).bits // 8 + if dtype_bytes == 0: + # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so + # use a fractional check expressed via bits. + col_bits = cols * DataType(dtype).bits + atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8 + if col_bits < atom_bits or col_bits % atom_bits != 0: + row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8 + atom_bytes = _swizzle_atom_bytes(swizzle_mode) + suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0) + raise ValueError( + f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces " + f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B " + f"swizzle atom selected by {swizzle_mode.name}. " + f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] " + f"to a multiple of " + f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements." + ) + else: + row_bytes = cols * dtype_bytes + atom_bytes = _swizzle_atom_bytes(swizzle_mode) + if row_bytes < atom_bytes or row_bytes % atom_bytes != 0: + suggestion = _suggest_swizzle_for_row_bytes(row_bytes) + min_cols = atom_bytes // dtype_bytes + raise ValueError( + f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces " + f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B " + f"swizzle atom selected by {swizzle_mode.name}. " + f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] " + f"to a multiple of {min_cols} elements (>= {atom_bytes}B at {dtype})." + ) + + # Atom rows is always 8 (see ``mma_atom_shape`` in tma_utils.py). + atom_rows = 8 + if rows < atom_rows or rows % atom_rows != 0: + raise ValueError( + f"alloc_mma shape={tuple(shape)} has shape[-2]={rows}, but the " + f"{swizzle_mode.name} atom requires shape[-2] to be a positive " + f"multiple of {atom_rows}. Use swizzle_mode='none', or widen shape[-2] " + f"to a multiple of {atom_rows}." + ) + + +# --------------------------------------------------------------------------- +# TMEMStages +# --------------------------------------------------------------------------- + + +def _meta_class(cls): + """Apply @meta_class decorator from ir_builder.""" + return _get_ir().meta_class(cls) + + +@_meta_class +class TMEMStages: + """Parse-time staged view over a TMEM buffer. + + Parameters + ---------- + buf : Buffer + The underlying TMEM buffer (e.g. f32 or f16 view). + col_start : int + First column of stage 0 in *buf*'s column space. + width : int + Number of columns per stage. + stages : int + Number of pipeline stages (default 1). + stride : int or None + Column distance between consecutive stages. When *None* (default), + equals *width* (stages are packed back-to-back). + """ + + def __init__(self, buf, col_start, width, stages=1, stride=None): + self.buf = buf + self.col_start = col_start + self.width = width + self.stages = stages + self.stride = width if stride is None else stride + + def _stage_base(self, stage): + return self.col_start + stage * self.stride + + def __getitem__(self, item): + if isinstance(item, tuple): + assert len(item) == 2, "TMEMStages expects region[stage] or region[stage, start:stop]" + stage, col_slice = item + assert isinstance(col_slice, slice), "TMEMStages tuple indexing requires a slice" + base = self._stage_base(stage) + start = 0 if col_slice.start is None else col_slice.start + stop = self.width if col_slice.stop is None else col_slice.stop + return self.buf[:, base + start : base + stop : col_slice.step] + base = self._stage_base(item) + return self.buf[:, base : base + self.width] + + +# --------------------------------------------------------------------------- +# TMEMPool +# --------------------------------------------------------------------------- + + +@_meta_class +class TMEMPool: + """Bump allocator over TMEM columns.""" + + def __init__( + self, + pool, + total_cols=512, + *, + cta_group=1, + alloc_warp=0, + dealloc_warp=None, + tmem_addr=None, + sync_after_alloc=True, + ): + # tcgen05 alloc/dealloc are warp-uniform PTX instructions: every lane + # in the chosen warp must participate, and exactly one warp in the + # CTA must execute them. The pool emits its own + # ``if warp_id() == target_warp: tcgen05.alloc(...)`` + # guard, using the cta->warp scope id ``T.warp_id()``. + # NOTE: synccheck currently false-deadlocks on kernels that declare a + # second warp-scope id (cpusim binds only one warp var); the generated + # CUDA is equivalent to ``thread_rank() // 32 == target_warp``. + self.pool = pool + self.total_cols = total_cols + self.cta_group = cta_group + self.alloc_warp = alloc_warp + self.dealloc_warp = alloc_warp if dealloc_warp is None else dealloc_warp + self.sync_after_alloc = sync_after_alloc + self.offset = 0 + self.max_offset = 0 + self._committed = False + self._deallocated = False + self._addr_buf = pool.alloc([1], "uint32", align=4) if tmem_addr is None else tmem_addr + + def _addr_slot(self): + try: + return self._addr_buf[0] + except TypeError: + return self._addr_buf + + @property + def addr(self): + return self._addr_slot() + + def _emit_warp_guard(self, target_warp, emit): + from tvm.script import tirx as T + + warp_id = T.warp_id() + with T.If(warp_id == target_warp): + with T.Then(): + emit() + + def _resolve_cols(self, shape, dtype, cols, layout=None): + if cols is not None: + return cols + bits = DataType(dtype).bits + if layout is not None: + # span("TCol") is in *element* (buffer dtype) units; one TMEM cell + # holds 32 bits regardless of the element type. + tcol_elems = int(layout.span("TCol")) + tcol_bits = tcol_elems * bits + assert tcol_bits % 32 == 0, ( + f"layout TCol span={tcol_elems} elems x {bits}b is not 32-bit aligned" + ) + return tcol_bits // 32 + assert len(shape) == 2, "TMEMPool.alloc() requires cols= for non-2D TMEM buffers" + total_bits = _shape_product(shape) * bits + rows = shape[0] + assert total_bits % (32 * rows) == 0, ( + f"Cannot infer TMEM columns from shape={shape}, dtype={dtype!r}; " + "please pass cols= explicitly" + ) + return total_bits // (32 * rows) + + def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None): + """Allocate a TMEM buffer. + + Parameters + ---------- + shape, dtype, cols + Standard buffer shape / dtype / column count. + layout + Explicit ``TileLayout``. Mutually exclusive with ``datapath``. + datapath : str | None + Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath, + ``"F"`` for M=64 non-``.ws`` scattered). When provided, the buffer's + layout is derived from ``tmem_datapath_layout(datapath, *shape)`` + so the row index reflects the *physical* TMEM lane occupation + (PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b`` + dispatches structurally check this layout to catch mismatched + atoms (e.g. a ``.16x*b`` M=128 read against a Layout F buffer). + Defaults to ``None``, which means Layout D's identity row→lane + mapping — keep this for shape ``(128, X)`` buffers that hold + an M=128 MMA accumulator. + """ + from tvm.tirx.layout import tmem_datapath_layout + + if layout is not None and datapath is not None: + raise ValueError("TMEMPool.alloc: pass at most one of layout= and datapath=") + if datapath is not None: + assert len(shape) == 2, "TMEMPool.alloc: datapath= requires a 2-D shape" + layout = tmem_datapath_layout(datapath, shape[0], shape[1]) + + ir = _get_ir() + cols = self._resolve_cols(shape, dtype, cols, layout) + col_start = self.offset + col_end = col_start + cols + assert col_end <= self.total_cols, f"TMEM overflow: {col_end} > {self.total_cols}" + if layout is None: + assert len(shape) == 2, "TMEMPool.alloc() requires layout= for non-2D TMEM buffers" + layout = _default_tmem_layout(shape[0], shape[1]) + res = ir.decl_buffer(shape, dtype, scope="tmem", allocated_addr=col_start, layout=layout) + self.offset = col_end + self.max_offset = max(self.max_offset, self.offset) + return res + + def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1): + """Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred layout. + + ``shape`` last two dims are ``(rows, SF_K * sf_reuse)`` (the last dim is + what gemm dispatch iterates over). When ``shape`` has 3 dims, the first + is treated as a pipe-depth outer. + """ + from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05 import sf_tmem_layout + + if len(shape) == 2: + pipe_depth, rows, last = None, shape[0], shape[1] + elif len(shape) == 3: + pipe_depth, rows, last = shape[0], shape[1], shape[2] + else: + raise ValueError( + f"alloc_sf expects 2D (rows, SF_K*sf_reuse) or 3D " + f"(pipe_depth, rows, SF_K*sf_reuse); got shape={shape}" + ) + assert last % sf_reuse == 0, ( + f"alloc_sf: shape last dim {last} must be divisible by sf_reuse={sf_reuse}" + ) + SF_K = last // sf_reuse + layout = sf_tmem_layout( + rows=rows, SF_K=SF_K, sf_per_mma=sf_per_mma, sf_reuse=sf_reuse, pipe_depth=pipe_depth + ) + return self.alloc(shape, dtype, layout=layout) + + def move_base_to(self, col): + self.offset = col + self.max_offset = max(self.max_offset, self.offset) + + def commit(self): + assert not self._committed, "TMEMPool.commit() can only be called once" + from tvm.script import tirx as T + + def emit_alloc(): + _emit_stmt( + T.ptx.tcgen05.alloc( + T.address_of(self.addr), n_cols=self.total_cols, cta_group=self.cta_group + ) + ) + if self.sync_after_alloc: + _emit_stmt(T.cuda.warp_sync()) + + self._emit_warp_guard(self.alloc_warp, emit_alloc) + self._committed = True + + def dealloc(self): + assert self._committed, "TMEMPool.dealloc() called before commit()" + assert not self._deallocated, "TMEMPool.dealloc() can only be called once" + self._deallocated = True + from tvm.script import tirx as T + + def emit_dealloc(): + _emit_stmt(T.ptx.tcgen05.relinquish_alloc_permit(cta_group=self.cta_group)) + _emit_stmt( + T.ptx.tcgen05.dealloc(self.addr, n_cols=self.total_cols, cta_group=self.cta_group) + ) + + self._emit_warp_guard(self.dealloc_warp, emit_dealloc) + + +# --------------------------------------------------------------------------- +# SMEMPool +# --------------------------------------------------------------------------- + + +@_meta_class +class SMEMPool: + """Bump allocator over a contiguous shared memory region. + + Parameters + ---------- + ptr : Var or None, optional + If omitted, an ``alloc_buffer([0], "uint8", scope="shared.dyn")`` is + created automatically and ``commit()`` must be called after all + allocations to emit the size annotation. + If a ``Var`` is provided, the caller manages the backing buffer and + ``commit()`` is a no-op. + """ + + def __init__(self, ptr=_POOL_UNSET): + ir = _get_ir() + if ptr is _POOL_UNSET: + self.buf = ir.alloc_buffer([0], "uint8", scope="shared.dyn") + self.ptr = self.buf.data + self._owns_buffer = True + else: + self.buf = None + self.ptr = ptr + self._owns_buffer = False + self.offset = 0 + self.max_offset = 0 + + def alloc( + self, + shape, + dtype="float32", + strides=None, + scope="shared.dyn", + align=0, + buffer_type="", + axis_separators=None, + layout="default", + ): + ir = _get_ir() + if align > 0: + self.offset = (self.offset + align - 1) // align * align + res = ir.decl_buffer( + shape, + dtype, + data=self.ptr, + strides=strides, + byte_offset=self.offset, + scope=scope, + align=align, + buffer_type=buffer_type, + axis_separators=axis_separators, + layout=layout, + ) + # Advance in bits then round up to bytes so sub-byte dtypes (e.g. + # float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it + # at 0 (bits // 8) and silently overlapping the next allocation. + self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8 + if self._owns_buffer: + self.max_offset = max(self.max_offset, self.offset) + return res + + def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024): + """Allocate MMA-compatible shared memory with an inferred swizzle layout.""" + from tvm.backend.cuda.operator.tile_primitive.tma_utils import ( + SwizzleMode, + mma_shared_layout, + ) + + if isinstance(swizzle_mode, str): + if swizzle_mode == "auto": + swizzle_mode = _auto_swizzle_mode(dtype) + elif swizzle_mode == "none": + swizzle_mode = SwizzleMode.SWIZZLE_NONE + else: + raise ValueError( + f"Unsupported swizzle_mode={swizzle_mode!r}; expected 'auto', 'none', " + "or SwizzleMode" + ) + _validate_mma_alloc_shape(shape, dtype, swizzle_mode) + layout = mma_shared_layout(dtype, swizzle_mode, shape) + return self.alloc(shape, dtype, align=align, layout=layout) + + def move_base_to(self, offset): + self.offset = offset + if self._owns_buffer: + self.max_offset = max(self.max_offset, self.offset) + + def commit(self, size=None): + """Emit pool size annotation into the IR. + + Must be called after all ``alloc()`` / ``move_base_to()`` calls. + + Parameters + ---------- + size : int, optional + Explicit shared memory size in bytes. When *None* (the default), + the high-water mark ``max_offset`` tracked by the allocator is used. + """ + if not self._owns_buffer: + return + ir = _get_ir() + frame_mod = _get_frame() + resolved = size if size is not None else self.max_offset + assert resolved >= self.max_offset, ( + f"Specified smem size ({resolved}) is smaller than " + f"the pool high-water mark ({self.max_offset})" + ) + attr_frame = ir.attr(self.ptr, "tirx.pool_max_bytes", resolved) + if isinstance(attr_frame, frame_mod.AttrFrame): + from functools import partial + + attr_frame.add_callback(partial(attr_frame.__exit__, None, None, None)) + attr_frame.__enter__() Review Comment:  In `SMEMPool.commit`, `attr_frame.add_callback` is called on `attr_frame` itself before it is entered. Since `attr_frame` is never exited manually, its callbacks are never run, leading to a frame leak/deadlock. It should be added to `parent_frame` instead. ```suggestion if isinstance(attr_frame, frame_mod.AttrFrame): from functools import partial parent_frame = frame_mod.current() parent_frame.add_callback(partial(attr_frame.__exit__, None, None, None)) attr_frame.__enter__() ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
