perf(kernel): optimize Blackwell grouped GEMM primitives

[spectrometerHBH]

Summary

• Add predicate support for SM100 block-scaled tcgen05 MMA and commit operations.
• Preserve and reuse hoisted shared-memory descriptors in grouped GEMM pipelines.
• Add configurable TMA L2 promotion and include it in tensor-map caching and encoding.
• Prune compile-time-only tile-scheduler branches and keep pipeline state warp-uniform.

Motivation

The downstream grouped FP8 GEMM needs predicated warp-specialized MMA issue, explicit TMA cache policy, and lower-overhead persistent scheduling. Without these changes, the generated SASS performs redundant descriptor and scheduler work on the critical producer paths.

Validation

• pre-commit run
• python -m pytest tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py -k test_copy_tma_codegen — 54 passed
• python -m pytest tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py -k test_gemm_smem_desc_hoist_vs_recompute — 2 passed
• python -m pytest tests/python/tirx/codegen/test_codegen_blackwell.py -k test_tcgen05_block_scaled_mma_predicate_codegen — 1 passed

[gemini-code-assist]

Code Review

This pull request refactors TVM's TIRx dialect, moving TilePrimitiveCall and DispatchContext to tvm.tirx.tile_primitive, replacing Predicate with LambdaExpr, and introducing robust dimension-surgery views (unflatten, flatten, select, narrow, sub, rearrange) to Buffer. It also extends PTX load/store intrinsics, optimizes tcgen05 block-scaled MMA and copy operations, and updates the benchmark utilities. The reviewer feedback highlights several critical issues: a scoping violation in tcgen05_cp.py from an over-simplified cache key, a missing cache_policy argument in ptx_ld_acquire, a missing cache_policy in the generated C signatures for relaxed and acquire loads, a potential AttributeError in scalar_dtype for raw Python scalars, fragile line-number parsing in monitor_gpu.sh, and import alias inconsistencies in binary.py.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

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.

Suggested change

	cache_key = f"smem_tmem_desc:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
	cache_key = f"smem_tmem_desc:{s_buf.name}:{int(ldo)}:{int(sdo)}:{int(swizzle)}"

[gemini-code-assist]

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.

Suggested change

	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,
	)
	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,
	)

[gemini-code-assist]

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.

Suggested change

	sig = (
	"(void* address, unsigned long long cache_policy)" if form == "weak" else "(void* address)"
	)
	sig = (
	"(void* address, unsigned long long cache_policy)"
	if form in ("weak", "relaxed", "acquire")
	else "(void* address)"
	)

[gemini-code-assist]

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.

Suggested change

	sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//'
	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

[gemini-code-assist]

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.

Suggested change

	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)
	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)

[gemini-code-assist]

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.

Suggested change

	from tvm.script import tirx as Tx
	from tvm.script import tirx as T

[gemini-code-assist]

Use T.max instead of Tx.max to align with the updated import alias T for tvm.script.tirx.

Suggested change

	return Tx.max(src_vals[0], src_vals[1])
	return T.max(src_vals[0], src_vals[1])