gemini-code-assist[bot] commented on code in PR #19988:
URL: https://github.com/apache/tvm/pull/19988#discussion_r3566002123
##########
python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py:
##########
@@ -372,21 +373,35 @@ def _canon_segment(iters):
# add_post_buffer_def_stmt.
# -----------------------------------------------------------------------------
def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
- cache_key =
f"smem_tmem_desc:{hash(s_buf)}:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
+ # Cache descriptor template at SMEM 0; patch addr per cp.
+ cache_key = f"smem_tmem_desc:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
Review Comment:
Removing `hash(s_buf)` from the cache key can lead to compilation errors due
to scoping violations when multiple distinct shared memory buffers are used in
the same kernel. Since `desc_buf` is allocated via
`add_post_buffer_def_stmt(s_buf, wrap)`, its scope is bound to `s_buf`'s
definition block. If another buffer `s_buf2` hits the cache and reuses
`desc_buf`, it will attempt to access `desc_buf` outside of its valid scope,
resulting in an undefined variable error during TVM compilation. Including
`s_buf.name` in the cache key (similar to `gemm_async/tcgen05.py`) resolves
this safely.
```suggestion
cache_key =
f"smem_tmem_desc:{s_buf.name}:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
```
##########
python/tvm/backend/cuda/op.py:
##########
@@ -3532,41 +3513,62 @@ def _resolve_cache_policy(cache_hint, cache_policy,
choices=_CP_ASYNC_BULK_CACHE
return const(0, dtype="uint64"), False
-def ptx_ld_acquire(addr, return_type, ptx_type, *, scope="gpu",
space="global"):
- """TVM intrinsic for scalar PTX ``ld.acquire.scope{.ss}.type`` loads.
-
- This wrapper covers the scalar no-cache-policy/no-vector instances of the
- PTX ISA ``ld.acquire`` form. ``scope``, state ``space``, PTX ``type`` and
- TVM ``return_type`` are explicit so callers can request either raw-bit or
- typed loads.
-
- Parameters
- ----------
- addr : PrimExpr
- The memory address to load.
-
- return_type : str
- TVM dtype returned by the load.
-
- ptx_type : str
- PTX type suffix such as ``"b32"``, ``"u64"``, or ``"s32"``.
-
- scope : str
- PTX memory scope: ``"cta"``, ``"cluster"``, ``"gpu"``, or ``"sys"``.
-
- space : str
- PTX state space suffix.
-
- Returns
- -------
- call : PrimExpr
- The loaded value.
- """
+def ptx_ld_acquire(
+ addr,
+ return_type,
+ ptx_type,
+ *,
+ scope="gpu",
+ space="global",
+ vec="",
+ dst=None,
+ cache_hint="",
+ cache_policy=None,
+ l1_evict="",
+ l2_evict="",
+ prefetch_size="",
+):
+ """TVM intrinsic for PTX ``ld.acquire.scope{.ss}...`` loads."""
_choice("scope", scope, _PTX_LD_SCOPE)
_choice("space", space, _PTX_LD_SPACE)
_choice("ptx_type", ptx_type, _PTX_LD_TYPE)
+ _choice("vec", vec, _PTX_LD_VEC)
+ cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint,
cache_policy)
+ to_dst = int(dst is not None)
+ if vec and dst is None:
+ raise ValueError("vec ld.acquire requires dst")
+ if to_dst:
+ return call_intrin(
+ "",
+ "tirx.ptx.ld_acquire",
+ dst,
+ addr,
+ cache_policy,
+ return_type,
+ scope,
+ space,
+ vec,
+ ptx_type,
+ int(has_cache_policy),
+ to_dst,
+ l1_evict,
+ l2_evict,
+ prefetch_size,
+ )
return call_intrin(
- return_type, "tirx.ptx.ld_acquire", addr, return_type, ptx_type,
scope, space
+ return_type,
+ "tirx.ptx.ld_acquire",
+ addr,
+ return_type,
+ scope,
+ space,
+ vec,
+ ptx_type,
+ int(has_cache_policy),
+ 0,
+ l1_evict,
+ l2_evict,
+ prefetch_size,
)
Review Comment:
In `ptx_ld_acquire` (when `to_dst` is False), `cache_policy` is completely
missing from the arguments passed to `call_intrin`. This is a major discrepancy
compared to `ptx_ld_relaxed` and `ptx_ld` which always pass `cache_policy` as
the second argument. This omission will cause compilation or runtime errors
because the generated C++ function expects `cache_policy` when `has_cache` is
true, but it is not passed.
```suggestion
return call_intrin(
return_type,
"tirx.ptx.ld_acquire",
addr,
cache_policy,
return_type,
scope,
space,
vec,
ptx_type,
int(has_cache_policy),
0,
l1_evict,
l2_evict,
prefetch_size,
)
```
##########
python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py:
##########
@@ -72,6 +83,19 @@ def compute_dtype_of(plan) -> str:
return widest
+def scalar_dtype(scalar) -> str:
+ dtype = getattr(scalar, "dtype", None)
+ if dtype is not None:
+ return str(dtype)
+ ty = getattr(scalar, "ty", None)
+ if ty is None and hasattr(scalar, "expr_ty"):
+ ty = scalar.expr_ty()
+ dtype = getattr(ty, "dtype", None)
+ if dtype is None:
+ raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing
PrimType")
+ return str(dtype)
Review Comment:
In `scalar_dtype`, if `scalar` is a raw Python `float` or `int`, the
function will raise an `AttributeError` because Python built-in types do not
have `dtype` or `ty` attributes. Adding a fallback check for
`isinstance(scalar, (float, int))` makes the function much more robust and
prevents unexpected crashes when raw Python scalars are passed.
```suggestion
def scalar_dtype(scalar) -> str:
if isinstance(scalar, float):
return "float32"
if isinstance(scalar, int):
return "int32"
dtype = getattr(scalar, "dtype", None)
if dtype is not None:
return str(dtype)
ty = getattr(scalar, "ty", None)
if ty is None and hasattr(scalar, "expr_ty"):
ty = scalar.expr_ty()
dtype = getattr(ty, "dtype", None)
if dtype is None:
raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing
PrimType")
return str(dtype)
```
##########
.agents/scripts/monitor_gpu.sh:
##########
@@ -23,18 +23,7 @@ while [[ $# -gt 0 ]]; do
--interval) INTERVAL="$2"; shift 2 ;;
--log) LOG="$2"; shift 2 ;;
-h|--help)
- cat <<'EOF'
-Watch a single GPU for foreign processes (anyone other than the current
-user) appearing during a long-running test. Intended companion to
-`/tir-test`: leave this running in a side terminal while pytest runs, and
-it will alert if someone else lands on the same GPU.
-
-Usage:
- monitor_gpu.sh # uses $CUDA_VISIBLE_DEVICES, defaults
to 0
- monitor_gpu.sh --gpu 3 # watch GPU 3
- monitor_gpu.sh --gpu 3 --interval 2 # poll every 2 seconds
- monitor_gpu.sh --log /tmp/gpu.log # also tee to a log file
-EOF
+ sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//'
Review Comment:
Hardcoding line numbers like `'2,12p'` in `sed` is extremely fragile because
any change to the top of the script (such as adding comments, changing imports,
or adding license headers) will break the help output. Keeping the heredoc `cat
<<'EOF'` is much more robust and standard.
```suggestion
cat <<'EOF'
Watch a single GPU for foreign processes (anyone other than the current
user) appearing during a long-running test. Intended companion to
'/tir-test': leave this running in a side terminal while pytest runs, and
it will alert if someone else lands on the same GPU.
Usage:
monitor_gpu.sh # uses $CUDA_VISIBLE_DEVICES,
defaults to 0
monitor_gpu.sh --gpu 3 # watch GPU 3
monitor_gpu.sh --gpu 3 --interval 2 # poll every 2 seconds
monitor_gpu.sh --log /tmp/gpu.log # also tee to a log file
EOF
```
##########
python/tvm/backend/cuda/operator/intrinsics/memory.py:
##########
@@ -144,158 +175,248 @@ def _validate_ld_space(space: str, allowed: set[str])
-> None:
)
-def _ptx_ld_helper_name(kind: str, return_dtype: str, ptx_type: str, scope:
str | None, space: str):
- parts = ["tvm_builtin_ptx_ld", kind]
- if scope is not None:
- parts.append(scope.replace("::", "_"))
- parts.extend([space.replace("::", "_"), ptx_type, return_dtype])
- return "_".join(parts)
+def _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size):
+ suffix = _cache_suffix("cache" if has_cache else "")
+ l1_evict = parse_str(l1_evict)
+ l2_evict = parse_str(l2_evict)
+ prefetch_size = parse_str(prefetch_size)
+ if l1_evict:
+ suffix += f".{l1_evict}"
+ if l2_evict:
+ suffix += f".{l2_evict}"
+ if prefetch_size:
+ suffix += f".{prefetch_size}"
+ return suffix
+
+
+def _ptx_shared_addr(space, ptr_name="address"):
+ if parse_str(space).startswith("shared"):
+ return (
+ f" unsigned int addr = (unsigned
int)__cvta_generic_to_shared({ptr_name});\n",
+ '"r"(addr)',
+ )
+ return "", f'"l"({ptr_name})'
+
+
+def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type):
+ if ptx_type == "u8" and vec_len == 1:
+ return " *reinterpret_cast<unsigned char*>(dst_ptr) =
static_cast<unsigned char>(r0);"
+ store_type = _PTX_VEC_STORE_TYPE[num_bytes]
+ if vec_len > 1:
+ return (
+ f" *reinterpret_cast<{store_type}*>(dst_ptr) = "
+ + "{"
+ + ", ".join(f"r{i}" for i in range(vec_len))
+ + "};"
+ )
+ return f" *reinterpret_cast<{store_type}*>(dst_ptr) = r0;"
-def _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache_hint):
- return_dtype, ptx_type, _scope, space, c_type = _parse_ld_attrs(
- return_dtype, ptx_type, None, space
+def _ptx_ld_form_parts(form, attr_args):
+ if form == "weak":
+ (
+ return_dtype,
+ weak,
+ space,
+ cop,
+ vec,
+ ptx_type,
+ has_cache_hint,
+ to_dst,
+ l1_evict,
+ l2_evict,
+ prefetch_size,
+ ) = attr_args
+ sem, scope = "", ""
+ elif form == "relaxed":
+ (
+ return_dtype,
+ scope,
+ space,
+ vec,
+ ptx_type,
+ has_cache_hint,
+ to_dst,
+ l1_evict,
+ l2_evict,
+ prefetch_size,
+ ) = attr_args
+ sem, weak, cop = "", False, ""
+ elif form == "acquire":
+ (
+ return_dtype,
+ scope,
+ space,
+ vec,
+ ptx_type,
+ has_cache_hint,
+ to_dst,
+ l1_evict,
+ l2_evict,
+ prefetch_size,
+ ) = attr_args
+ sem, weak, cop = "", False, ""
+ elif form == "volatile":
+ return_dtype, space, vec, ptx_type, to_dst, prefetch_size = attr_args
+ sem, scope, weak, cop = "", "", False, ""
+ has_cache_hint, l1_evict, l2_evict = False, "", ""
+ elif form == "mmio":
+ return_dtype, sem, scope, space, ptx_type, to_dst = attr_args
+ weak, cop, vec = False, "", ""
+ has_cache_hint, l1_evict, l2_evict, prefetch_size = False, "", "", ""
+ else:
+ raise ValueError(f"unknown ld form {form!r}")
+
+ return_dtype, ptx_type, scope, space, c_type, constraint = _parse_ld_attrs(
+ return_dtype, ptx_type, scope if form in ("relaxed", "acquire") else
None, space
)
+ sem = parse_str(sem)
+ scope = parse_str(scope)
+ space = parse_str(space)
cop = parse_str(cop)
- if cop not in _PTX_LD_COPS:
+ vec = parse_str(vec)
+ l1_evict = parse_str(l1_evict)
+ l2_evict = parse_str(l2_evict)
+ prefetch_size = parse_str(prefetch_size)
+ weak = _bool_attr(weak)
+ has_cache = _bool_attr(has_cache_hint)
+ to_dst = _bool_attr(to_dst)
+ if cop and cop not in _PTX_LD_COPS:
raise ValueError(f"Unsupported PTX ld cache operation {cop!r}")
- weak = bool(int(weak)) if hasattr(weak, "value") else bool(weak)
- has_cache = (
- bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else
bool(has_cache_hint)
+ if vec and vec not in _PTX_VEC:
+ raise ValueError(f"Unsupported PTX ld vector modifier {vec!r}")
+ if l1_evict and l1_evict not in _PTX_L1_EVICT:
+ raise ValueError(f"Unsupported PTX ld L1 eviction {l1_evict!r}")
+ if l2_evict and l2_evict not in _PTX_L2_EVICT:
+ raise ValueError(f"Unsupported PTX ld L2 eviction {l2_evict!r}")
+ if prefetch_size and prefetch_size not in _PTX_PREFETCH:
+ raise ValueError(f"Unsupported PTX ld prefetch size {prefetch_size!r}")
+ if form == "mmio":
+ if sem not in ("acquire", "relaxed") or scope != "sys" or space !=
"global":
+ raise ValueError("ld.mmio requires sem in {acquire, relaxed},
scope=sys, space=global")
+ prefix = f"ld.mmio.{sem}.{scope}"
+ elif form == "relaxed":
+ if not scope:
+ raise ValueError("ld.relaxed requires scope")
+ _validate_ld_space(space, _PTX_LD_SPACES)
+ prefix = f"ld.relaxed.{scope}{_dot(space)}"
+ elif form == "acquire":
+ if not scope:
+ raise ValueError("ld.acquire requires scope")
+ _validate_ld_space(space, _PTX_LD_SPACES)
+ prefix = f"ld.acquire.{scope}{_dot(space)}"
+ elif form == "volatile":
+ _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES)
+ prefix = f"ld.volatile{_dot(space)}"
+ else:
+ _validate_ld_space(space, _PTX_LD_WEAK_SPACES)
+ prefix = f"ld{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}"
+ level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size)
+ vec_len = int(vec[1:]) if vec else 1
+ if vec and not to_dst:
+ raise ValueError("vector ld requires to_dst")
+ elem_bytes = (
+ 8
+ if ptx_type.endswith("64")
+ else 2
+ if ptx_type in ("u16", "s16", "b16")
+ else 1
+ if ptx_type in ("u8", "s8", "b8")
+ else 4
)
- _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES | {"param::entry",
"param::func"})
- spec = _PTX_LD_TYPES[ptx_type]["constraint"]
- addr_decl = ""
- addr_operand = '"l"(address)'
- if space.startswith("shared"):
- addr_decl = " unsigned int addr = (unsigned
int)__cvta_generic_to_shared(address);\n"
- addr_operand = '"r"(addr)'
- modifiers = f"{'.weak' if weak else ''}.{space}{('.' + cop) if cop else
''}"
- cache_inst = ".L2::cache_hint" if has_cache else ""
- cache_slot = ", %2" if has_cache else ""
- cache_operand = ', "l"(cache_policy)' if has_cache else ""
- name = (
- "tvm_builtin_ptx_ld"
- f"{'_weak' if weak else ''}_{space.replace('::', '_').replace('.',
'_')}"
- f"{('_' + cop) if cop else ''}_{ptx_type}_{return_dtype}"
- f"{'_cache_hint' if has_cache else ''}"
+ num_bytes = vec_len * elem_bytes if vec else elem_bytes
+ name_parts = [
+ "tvm_builtin_ptx_ld",
+ form if form != "weak" else ("weak" if weak else "plain"),
+ ]
+ if sem:
+ name_parts.append(_safe_attr(sem))
+ if scope:
+ name_parts.append(_safe_attr(scope))
+ name_parts.extend(
+ [
+ _safe_attr(space),
+ _safe_attr(cop) if cop else "",
+ _safe_attr(vec) if vec else "",
+ ptx_type,
+ return_dtype if not to_dst else "to_dst",
+ ]
)
+ if has_cache:
+ name_parts.append("cache_hint")
+ if l1_evict:
+ name_parts.append(_safe_attr(l1_evict))
+ if l2_evict:
+ name_parts.append(_safe_attr(l2_evict))
+ if prefetch_size:
+ name_parts.append(_safe_attr(prefetch_size))
+ name = "_".join(p for p in name_parts if p)
+ cache_operand = ', "l"(cache_policy)' if has_cache else ""
+ addr_decl, addr_operand = _ptx_shared_addr(space, "src_ptr" if to_dst else
"address")
+ if to_dst:
+ reg_decls = "".join(f" {c_type} r{i};\n" for i in range(vec_len))
+ if vec_len > 1:
+ out_slot = "{" + ", ".join(f"%{i}" for i in range(vec_len)) + "}"
+ out_constraints = ", ".join(f'"={constraint}"(r{i})' for i in
range(vec_len))
+ addr_idx = vec_len
+ else:
+ out_slot = "%0"
+ out_constraints = f'"={constraint}"(r0)'
+ addr_idx = 1
+ cache_slot = f", %{addr_idx + 1}" if has_cache else ""
+ instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
+ body = (
+ f"{addr_decl}{reg_decls}"
+ f' asm volatile("{instr} {out_slot},
[%{addr_idx}]{cache_slot};"\n'
+ f" : {out_constraints}\n"
+ f" : {addr_operand}{cache_operand});\n"
+ f"{_ptx_ld_vec_store(num_bytes, vec_len, ptx_type)}"
+ )
+ return (
+ name,
+ "(void* dst_ptr, void* src_ptr, unsigned long long cache_policy)",
+ "void",
+ "",
+ body,
+ )
+ cache_slot = ", %2" if has_cache else ""
+ instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
body = (
f" {c_type} ret;\n"
f"{addr_decl}"
- f' asm volatile("ld{modifiers}{cache_inst}.{ptx_type} %0,
[%1]{cache_slot};" '
- f': "={spec}"(ret) : {addr_operand}{cache_operand});\n'
+ f' asm volatile("{instr} %0, [%1]{cache_slot};"\n'
+ f' : "={constraint}"(ret)\n'
+ f" : {addr_operand}{cache_operand});\n"
" return ret;"
)
- return name, c_type, return_dtype, body
-
-
-device_intrinsic(
- "ptx_ld",
- n_attrs=6,
- helper_name=lambda _addr, _cache_policy, return_dtype, weak, space, cop,
ptx_type, has_cache: (
- _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[0]
- ),
- c_signature="(void* address, unsigned long long cache_policy)",
- return_type=lambda _addr, _cache_policy, return_dtype, weak, space, cop,
ptx_type, has_cache: (
- _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[1]
- ),
- tvm_return_type=lambda _addr,
- _cache_policy,
- return_dtype,
- _weak,
- _space,
- _cop,
- _ptx_type,
- _has_cache: (parse_str(return_dtype)),
- body=lambda _addr, _cache_policy, return_dtype, weak, space, cop,
ptx_type, has_cache: (
- _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[3]
- ),
-)
-
-
-def _ptx_ld_acquire_parts(return_dtype, ptx_type, scope, space):
- return_dtype, ptx_type, scope, space, c_type = _parse_ld_attrs(
- return_dtype, ptx_type, scope, space
+ sig = (
+ "(void* address, unsigned long long cache_policy)" if form == "weak"
else "(void* address)"
)
Review Comment:
In `_ptx_ld_form_parts`, the C signature `sig` only includes `cache_policy`
if `form == "weak"`. However, `relaxed` and `acquire` forms also support and
use `cache_policy` when `has_cache` is true. If `has_cache` is true for
`relaxed` or `acquire`, the generated C function will fail to compile because
`cache_policy` is referenced in the `asm` block but is not present in the
function signature. Changing the condition to `if form in ("weak", "relaxed",
"acquire")` resolves this.
```suggestion
sig = (
"(void* address, unsigned long long cache_policy)"
if form in ("weak", "relaxed", "acquire")
else "(void* address)"
)
```
##########
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py:
##########
@@ -100,6 +103,10 @@ def _compute_fdiv(src_vals, extras, dt):
return src_vals[0] / src_vals[1]
+def _compute_maximum(src_vals, extras, dt):
+ return Tx.max(src_vals[0], src_vals[1])
Review Comment:
Use `T.max` instead of `Tx.max` to align with the updated import alias `T`
for `tvm.script.tirx`.
```suggestion
return T.max(src_vals[0], src_vals[1])
```
##########
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py:
##########
@@ -31,12 +33,13 @@
import operator
from typing import Any
+from tvm.script import tirx as Tx
Review Comment:
To maintain consistency with the rest of the codebase where `T` is used for
`tvm.script.tirx` and `Tx` is used for `tvm.script.tirx.tile`, please import
`tirx` as `T`.
```suggestion
from tvm.script import tirx as T
```
--
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]
