mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-29 07:53:03 +08:00
Zhang JasonandLeiWang1999
3d95d65000
* [AMD][RDNA4] Fix RDNA4 (gfx1201 / Wave32) CI failures
- reduce.h: Replace hardcoded wave64 assumptions with runtime
__builtin_amdgcn_wavefrontsize() in SharedReduceWarp and AllReduce
- scan.h: Add SEG=32 support in InclusiveScan/CumSum/CumMax for RDNA wave32
- finalize_reducer.cc: Use TargetGetWarpSize() instead of hardcoded 64
- carver/arch/__init__.py + rdna.py: Add gfx12 (RDNA4) support alongside gfx11
- wmma_macro_generator.py: Fix k_pack indexing in ldmatrix_a/b for k_pack > 1
These changes correspond to PR #2210 (Fix RDNA4 gfx1201/Wave32 CI Failures),
rebased cleanly onto upstream main (a8d93798).
* [AMD][HIP] Fix HIPMath dispatch for vector dtype in T.Parallel
When BLOCK_M > threads, TileLang generates vector types (e.g. float32x4)
for T.Parallel loops. HIPMath::operator() only handled scalar types, so
vector dtypes fell through to return causing Unresolved call errors
in codegen for math intrinsics like tirx.exp2.
Fix: add a vector-type early-return that recurses on element_of(), so
float32xN correctly resolves to exp2f the same as scalar float32.
This complements the tirx.* registration fixes in #2282 and enables
Online Softmax on gfx1151 with BM > threads (e.g. BM=1024, TH=256).
* [Review] Fix 3 issues found in code review
1. tilelang/carver/arch/__init__.py: error message was not updated
when gfx12 support was added — still said gfx11 only. Now says
gfx11/gfx12 targets only, consistent with rdna.py.
2. src/tl_templates/hip/reduce.h: the static_assert(Threads % kWarpSize == 0)
was removed when kWarpSize became a runtime value, leaving no safety
check. Added equivalent runtime assert().
3. src/tl_templates/hip/scan.h: run_auto() else branch silently fell
through to SEG=64 for any non-32 wavefront size. Added assert to
catch unsupported wavefront sizes early.
* update codes for format checking
* [Review] Address review feedback and rebase on upstream main
- intrin_rule_hip.cc: remove redundant HIPMath scalar recursion; upstream
PR #2286 already handles vector dtypes via DispatchPureExternScalarized
at the IR level (correct for HIP packed carrier types)
- wmma_macro_generator.py: rename l/r to row_off/col_off to fix Ruff E741
ambiguous variable name warnings
- scan.h: guard run_auto() branches with threads>=32 / threads>=64 so a
threads==32 kernel on Wave64 doesn't compute TILE=0 and divide by zero
- reduce.h: replace assert (stripped in Release) with unconditional early
return when Threads % kWarpSize != 0, preventing silent num_warps==0
* update codes for format checking
* fix
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
147 lines
4.9 KiB
Python
147 lines
4.9 KiB
Python
import pytest
|
|
|
|
|
|
import tilelang.testing
|
|
from tilelang import tvm as tvm
|
|
from tvm.target import Target
|
|
|
|
import tilelang.utils.target as target_utils
|
|
from tilelang.utils.target import (
|
|
determine_target,
|
|
normalize_rocm_arch,
|
|
rocm_warp_size_for_arch,
|
|
target_get_mcpu,
|
|
target_get_rdna_generation,
|
|
target_get_warp_size,
|
|
target_is_cdna,
|
|
target_is_rdna,
|
|
)
|
|
|
|
|
|
def _hip_target(mcpu: str) -> Target:
|
|
return Target({"kind": "hip", "mcpu": mcpu})
|
|
|
|
|
|
def test_normalize_rocm_arch_strips_feature_suffix():
|
|
assert normalize_rocm_arch("gfx1151:sramecc+:xnack-") == "gfx1151"
|
|
assert normalize_rocm_arch("gfx942") == "gfx942"
|
|
assert normalize_rocm_arch("") is None
|
|
assert normalize_rocm_arch("sm_90") is None
|
|
assert rocm_warp_size_for_arch("gfx1151") == 32
|
|
assert rocm_warp_size_for_arch("gfx1030") == 32
|
|
assert rocm_warp_size_for_arch("gfx1200") == 32
|
|
assert rocm_warp_size_for_arch("gfx942") == 64
|
|
|
|
|
|
def test_target_mcpu_helpers():
|
|
target = _hip_target("gfx1151:sramecc+:xnack-")
|
|
assert target_get_mcpu(target) == "gfx1151"
|
|
|
|
|
|
def test_determine_target_adds_rdna_thread_warp_size():
|
|
target = determine_target({"kind": "hip", "mcpu": "gfx1151"}, return_object=True)
|
|
assert target_get_mcpu(target) == "gfx1151"
|
|
assert int(target.attrs["thread_warp_size"]) == 32
|
|
|
|
|
|
def test_determine_target_adds_known_gfx12_thread_warp_size():
|
|
target = determine_target({"kind": "hip", "mcpu": "gfx1200"}, return_object=True)
|
|
assert target_get_mcpu(target) == "gfx1200"
|
|
assert int(target.attrs["thread_warp_size"]) == 32
|
|
|
|
|
|
def test_determine_target_rejects_legacy_option_string():
|
|
with pytest.raises(AssertionError, match="Pass target options as a dict"):
|
|
determine_target("hip -mcpu=gfx1151", return_object=True)
|
|
|
|
|
|
def test_auto_target_prefers_rocm_pytorch_over_cuda_toolkit(monkeypatch):
|
|
monkeypatch.setattr(target_utils.torch.version, "hip", "test", raising=False)
|
|
monkeypatch.setattr(target_utils, "check_hip_availability", lambda: True)
|
|
monkeypatch.setattr(target_utils, "check_cuda_availability", lambda: True)
|
|
monkeypatch.setattr(target_utils, "_detect_torch_rocm_arch", lambda: "gfx1151")
|
|
|
|
target = determine_target("auto", return_object=True)
|
|
assert target.kind.name == "hip"
|
|
assert target_get_mcpu(target) == "gfx1151"
|
|
assert int(target.attrs["thread_warp_size"]) == 32
|
|
|
|
|
|
def test_rdna_gfx1151_target_classification():
|
|
target = _hip_target("gfx1151")
|
|
assert target_is_rdna(target)
|
|
assert not target_is_cdna(target)
|
|
assert target_get_rdna_generation(target) == 11
|
|
assert target_get_warp_size(target) == 32
|
|
|
|
|
|
def test_carver_routes_rdna_without_instantiating_device(monkeypatch):
|
|
import torch
|
|
|
|
monkeypatch.setattr(torch.version, "hip", None, raising=False)
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
|
if hasattr(torch, "mps"):
|
|
monkeypatch.setattr(torch.mps, "is_available", lambda: False, raising=False)
|
|
|
|
import tilelang.carver.arch as arch_mod
|
|
|
|
def fake_rdna(target):
|
|
return ("rdna", target)
|
|
|
|
monkeypatch.setattr(arch_mod, "RDNA", fake_rdna)
|
|
arch = arch_mod.get_arch(_hip_target("gfx1151"))
|
|
assert arch[0] == "rdna"
|
|
assert target_get_mcpu(arch[1]) == "gfx1151"
|
|
|
|
|
|
@tilelang.testing.requires_rocm
|
|
def test_carver_rejects_unsupported_rdna_generations(monkeypatch):
|
|
import tilelang.carver.arch as arch_mod
|
|
|
|
def fake_cdna(target):
|
|
return ("cdna", target)
|
|
|
|
monkeypatch.setattr(arch_mod, "CDNA", fake_cdna)
|
|
with pytest.raises(ValueError, match="gfx11 targets only"):
|
|
arch_mod.get_arch(_hip_target("gfx1200"))
|
|
|
|
|
|
@tilelang.testing.requires_rocm
|
|
def test_rdna_device_model_rejects_gfx12_before_device_probe():
|
|
from tilelang.carver.arch.rdna import RDNA
|
|
|
|
with pytest.raises(ValueError, match="gfx11 targets only"):
|
|
RDNA(_hip_target("gfx1200"))
|
|
|
|
|
|
@tilelang.testing.requires_rocm
|
|
def test_rdna_tensor_instruction_lookup_is_generation_aware():
|
|
from tilelang.carver.arch.rdna import RDNA
|
|
|
|
arch = RDNA.__new__(RDNA)
|
|
arch.rdna_generation = 11
|
|
assert arch.get_avaliable_tensorintrin_shapes() == [[16, 16]]
|
|
assert isinstance(arch.available_tensor_instructions, list)
|
|
|
|
arch.rdna_generation = 12
|
|
with pytest.raises(ValueError, match="Unsupported RDNA generation"):
|
|
arch.get_avaliable_tensorintrin_shapes()
|
|
|
|
|
|
def test_rdna_internal_tuning_config_allows_mcpu_overrides():
|
|
from tilelang.carver.arch.rdna import _get_rdna_tuning_config
|
|
|
|
default = _get_rdna_tuning_config("gfx1100")
|
|
gfx1151 = _get_rdna_tuning_config("gfx1151:sramecc+:xnack-")
|
|
|
|
assert default.preferred_warps_per_block == 4
|
|
assert default.pipeline_stage == 1
|
|
assert default.reduction_step_for_dtype_bits(16) == 32
|
|
assert gfx1151.preferred_warps_per_block == 4
|
|
assert gfx1151.pipeline_stage == 2
|
|
assert gfx1151.reduction_step_for_dtype_bits(16) == 32
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|