#!/usr/bin/env python3 # Copyright Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT """Test basic functionality of rocRoller's GEMM client.""" import contextlib import functools import itertools import os import shutil import subprocess from dataclasses import dataclass from pathlib import Path import pytest import yaml SOLUTION_NOT_SUPPORTED_ON_ARCH = 3 build = Path(__file__).parent.parent / "build" if os.getenv("ROCROLLER_BUILD_DIR") is not None: build = Path(os.getenv("ROCROLLER_BUILD_DIR")) gemm = (build / "client" / "rocroller-gemm").resolve() # Python 3.11 has contextlib.chdir but 3.10 doesn't @contextlib.contextmanager def chdir(directory): current_directory = os.getcwd() try: os.chdir(directory) yield finally: os.chdir(current_directory) # # Helpers # @functools.cache def rocm_gfx(): """Return GPU architecture (gfxXXXX) for local GPU device.""" output = None try: output = subprocess.run( ["rocminfo"], capture_output=True, text=True, check=True ).stdout except subprocess.CalledProcessError: return None for line in output.splitlines(): if line.startswith(" Name:"): _, arch, *_ = list(map(lambda x: x.strip(), line.split())) if arch.startswith("gfx"): return arch return None def match_architecture(solution_params): """Look for '--arch' option and check if it matches the local device. Returns True if either: 1. No --arch flag present. 2. The --arch flag matches the local device. """ if not ("--arch" in solution_params): return True arch_index = solution_params.index("--arch") + 1 arch = solution_params[arch_index] return arch == rocm_gfx() def check_returncode(p): """Checks return code of GEMM client. Returns True if the GEMM client returned SOLUTION_NOT_SUPPORTED_ON_ARCH. Raises an error if the GEMM client returned any other non-zero code. This usually means a correctness failure. Returns False if the GEMM client returned 0 (ie, OK). """ if p.returncode != 0: if p.returncode == SOLUTION_NOT_SUPPORTED_ON_ARCH: return True else: raise RuntimeError("Client failure: correctness or general failure") return False def write_solution_config_if_present(tmp_path, solution_params): """Write contents of '--config' option to .yaml configuration file (if present). This helper looks for '--config' in the solution params. It then takes the text of the following argument and writes it to a file named 'config.yaml', and replaces the argument. """ if not ("--config" in solution_params): return solution_params solution_params = list(solution_params) # copy config_index = solution_params.index("--config") + 1 config = solution_params[config_index] config_path = tmp_path / "config.yaml" config_path.write_text(config) solution_params[config_index] = config_path return solution_params # # Solution parameter helpers # @dataclass class Types: """Container for GEMM type parameters.""" A: str B: str C: str D: str def client_arguments(self): return [ "--type_A", self.A, "--type_B", self.B, "--type_C", self.C, "--type_D", self.D, ] @dataclass class Scale: """Container for MX GEMM scaling parameters.""" argument: str # A or B mode: str # Separate, SingleScale, etc path: bool # load through LDS value: float # for SingleScale, the value blockSize: int # for scale block size scaleType: str # data type of the scale values def client_arguments(self): params = [] if self.mode is not None: params.extend(["--scale_" + self.argument, self.mode]) if self.value is not None: params.extend(["--scaleValue_" + self.argument, str(self.value)]) if self.path is not None: params.extend(["--loadScale_" + self.argument, self.path]) if self.mode == "Separate" or self.mode == "SingleScale": if self.scaleType is not None: params.extend(["--scaleType_" + self.argument, self.scaleType]) else: params.extend(["--scaleType_" + self.argument, "E8M0"]) return params def maybe_add_block_size(self, params): if self.blockSize is None: return scaleBlockSizeStr = "--scaleBlockSize" if scaleBlockSizeStr not in params: params.extend([scaleBlockSizeStr, str(self.blockSize)]) @dataclass class Prefetch: """Container for GEMM prefetching parameters.""" in_flight: int lds_factor: int def client_arguments(self): params = [] if self.in_flight != 0: params.extend( [ "--prefetch", "--prefetchInFlight", str(self.in_flight), "--prefetchLDSFactor", str(self.lds_factor), ] ) return params DP_GEMM = """\ --- architecture: ArchString: gfxunknown Xnack: false Sramecc: false AsicRevisionId: -1 mac_m: 64 mac_n: 64 mac_k: 64 wave_m: 32 wave_n: 32 wave_k: 2 wave_b: 1 workgroup_size_x: 128 workgroup_size_y: 2 workgroupMappingDim: -1 workgroupRemapXCC: false workgroupRemapXCCValue: -1 workgroup_cluster_size_x: 0 workgroup_cluster_size_y: 0 workgroup_cluster_size_z: 0 load_A: BufferToLDSViaVGPR load_B: BufferToLDSViaVGPR padLDS_A: [0, 0] padLDS_B: [0, 0] store: VGPRToGlobalMemoryViaLDSWithBuffer prefetch: false prefetchInFlight: 0 prefetchLDSFactor: 0 prefetchMixMemOps: false betaInFma: true scheduler: Priority schedulerCost: LinearWeighted types: trans_A: N trans_B: N type_A: float type_B: float type_C: float type_D: float type_acc: float scale_A: None scaleType_A: None scale_B: None scaleType_B: None scaleBlockSize: 0 scalePreTileA: [] scalePreTileB: [] scaleShuffleTileA: [] scaleShuffleTileB: [] scaleSkipPermlane: None pretileA: [] pretileB: [] tailLoops: true streamK: None loadScale_A: BufferToVGPR loadScale_B: BufferToVGPR swizzleScale: false swizzleTileSize: m: 0 k: 0 n: 0 l: 0 prefetchScale: false ... """ DP_HGEMM = """\ --- architecture: ArchString: gfx90a Xnack: false Sramecc: false AsicRevisionId: -1 mac_m: 64 mac_n: 64 mac_k: 64 wave_m: 32 wave_n: 32 wave_k: 8 wave_b: 1 workgroup_size_x: 128 workgroup_size_y: 2 workgroupMappingDim: -1 workgroupRemapXCC: false workgroupRemapXCCValue: -1 workgroup_cluster_size_x: 0 workgroup_cluster_size_y: 0 workgroup_cluster_size_z: 0 load_A: BufferToLDSViaVGPR load_B: BufferToLDSViaVGPR padLDS_A: [0, 0] padLDS_B: [0, 0] store: VGPRToGlobalMemoryViaLDSWithBuffer prefetch: false prefetchInFlight: 0 prefetchLDSFactor: 0 prefetchMixMemOps: false betaInFma: true scheduler: Priority schedulerCost: LinearWeighted tailLoops: true types: trans_A: N trans_B: N type_A: half type_B: half type_C: half type_D: half type_acc: float scale_A: None scaleType_A: None scale_B: None scaleType_B: None scaleBlockSize: 0 scalePreTileA: [] scalePreTileB: [] scaleShuffleTileA: [] scaleShuffleTileB: [] scaleSkipPermlane: None pretileA: [] pretileB: [] loadScale_A: BufferToVGPR loadScale_B: BufferToVGPR swizzleScale: false swizzleTileSize: m: 0 k: 0 n: 0 l: 0 prefetchScale: false streamK: None ... """ DP_HGEMM_GFX120X = """\ --- architecture: ArchString: gfx1201 Xnack: false Sramecc: false AsicRevisionId: -1 mac_m: 64 mac_n: 64 mac_k: 64 wave_m: 16 wave_n: 16 wave_k: 16 wave_b: 1 workgroup_size_x: 64 workgroup_size_y: 2 workgroupMappingDim: -1 workgroupRemapXCC: false workgroupRemapXCCValue: -1 workgroup_cluster_size_x: 0 workgroup_cluster_size_y: 0 workgroup_cluster_size_z: 0 load_A: BufferToLDSViaVGPR load_B: BufferToLDSViaVGPR padLDS_A: [0, 0] padLDS_B: [0, 0] store: VGPRToGlobalMemoryViaLDSWithBuffer prefetch: false prefetchInFlight: 0 prefetchLDSFactor: 0 prefetchMixMemOps: false betaInFma: true scheduler: Priority schedulerCost: LinearWeighted tailLoops: true types: trans_A: N trans_B: N type_A: half type_B: half type_C: half type_D: half type_acc: float scale_A: None scaleType_A: None scale_B: None scaleType_B: None scaleBlockSize: 0 scalePreTileA: [] scalePreTileB: [] scaleShuffleTileA: [] scaleShuffleTileB: [] scaleSkipPermlane: None pretileA: [] pretileB: [] loadScale_A: BufferToVGPR loadScale_B: BufferToVGPR swizzleScale: false swizzleTileSize: m: 0 k: 0 n: 0 l: 0 prefetchScale: false streamK: None ... """ def type_configurations(): """Return list of type combinations to test.""" typeAs = ["fp4"] typeBs = ["fp4", "fp8"] typeDs = ["half", "float"] return [Types(A, B, D, D) for A, B, D in itertools.product(typeAs, typeBs, typeDs)] def scale_configurations(argument): """Return list of MX scale modes to test for each of A and B.""" modes = [None, "None", "Separate", "SingleScale"] paths = ["BufferToVGPR", "BufferToLDSViaVGPR"] values = [0.5, 1.0] blockSize = 32 scaleType = "E8M0" rv = [] for mode in modes: if mode is not None and mode == "Separate": rv.extend( [ Scale(argument, mode, path, None, blockSize, scaleType) for path in paths ] ) elif mode is not None and mode == "SingleScale": rv.extend( [ Scale(argument, mode, None, value, None, scaleType) for value in values ] ) else: rv.append(Scale(argument, mode, None, None, None, None)) return rv def prefetch_configurations(): """Return list of prefetching modes to test.""" rv = [Prefetch(0, 0)] for lds_factor in [0, 2]: rv.append(Prefetch(2, lds_factor)) return rv def build_solution_params(): """Build giant list of solution parameter combinations to test.""" solution_params = [ # data-parallel gemm, float, params from command line [], # data-parallel gemm, float, params from config file ["--config", DP_GEMM], # streamk gemm, float, params from command line ["--streamK", "Standard"], ] for type, prefetch, scaleA, scaleB in itertools.product( type_configurations(), prefetch_configurations(), scale_configurations("A"), scale_configurations("B"), ): # XXXX Mixing and outputting to half precision fails correctness checks. if (type.A != type.B) and (type.D == "half"): continue params = [ "--arch", "gfx950", "--mac_M", "256", "--mac_N", "256", "--mac_K", "64", "--mi", "32x32x64x1", ] for x in [type, prefetch, scaleA, scaleB]: params.extend(x.client_arguments()) scaleA.maybe_add_block_size(params) scaleB.maybe_add_block_size(params) if scaleA.mode == "Separate" or scaleB.mode == "Separate": params.extend(["--sts", "64x4/64x4"]) solution_params.append(params) return solution_params def build_problem_params(): """Build list of problem parameters to test.""" # Should consider making this a container return [["--m", "512", "--n", "512", "--k", "256", "--numWGs", "4"]] def build_wgm_params(): """Build list of workgroup-mapping parameters to test.""" size = [1024, 1024, 512] tile_size = [64, 64, 64] num_tiles = [size[i] // tile_size[i] for i in range(len(size))] params = [] for dimension in [0, 1]: # the last entry will have no "tail block" for wgm in [1, 3, 5, 16, num_tiles[dimension]]: solution_params = [ f"--mac_m={tile_size[0]}", f"--mac_n={tile_size[1]}", f"--mac_k={tile_size[2]}", f"--workgroupMappingDim={dimension}", f"--workgroupMappingValue={wgm}", ] problem_params = [f"--m={size[0]}", f"--n={size[1]}", f"--k={size[2]}"] params.append([solution_params, problem_params]) return params # # GEMM 'generate' and 'validate' helpers # def gemm_validate_single_stage(tmp_path, solution_params, problem_params): solution_params = write_solution_config_if_present(tmp_path, solution_params) if not match_architecture(solution_params): return cmd = [gemm] cmd.extend(["generate"]) cmd.extend(solution_params) cmd.extend(["validate"]) cmd.extend(problem_params) check_returncode(subprocess.run(cmd)) def gemm_validate_two_stage_codeobject(tmp_path, solution_params, problem_params): solution_params = write_solution_config_if_present(tmp_path, solution_params) co = tmp_path / "test.co" cmd = [gemm] cmd.extend(["generate", "--co", co]) cmd.extend(solution_params) skip = check_returncode(subprocess.run(cmd)) if skip: return if not match_architecture(solution_params): return cmd = [gemm] cmd.extend(["validate", "--load", co]) cmd.extend(problem_params) subprocess.run(cmd, check=True) def gemm_validate_two_stage_assembly(tmp_path, solution_params, problem_params): solution_params = write_solution_config_if_present(tmp_path, solution_params) asm = tmp_path / "test.s" # get these working; load command from .co cmd = [gemm] cmd.extend(["generate", "--asm", asm]) cmd.extend(solution_params) skip = check_returncode(subprocess.run(cmd)) if skip: return if not match_architecture(solution_params): return cmd = [gemm] cmd.extend(["validate", "--load", asm]) cmd.extend(problem_params) subprocess.run(cmd, check=True) # # PyTest tests! # def test_gemm_example(tmp_path): """GEMM 'example' subcommand.""" # "gemm example" with no output file should fail with pytest.raises(subprocess.CalledProcessError): subprocess.run([gemm, "example"], check=True) # "gemm example example.yaml" should write to example.yaml example = tmp_path / "example.yaml" subprocess.run([gemm, "example", example], check=True) assert example.exists() def test_gemm_options(tmp_path): """GEMM options.""" example = tmp_path / "example.yaml" example_problem = tmp_path / "example_problem.yaml" def run_and_load_example_yaml(cmd): subprocess.run(cmd, check=True) yaml_contents = example.read_text() return yaml.load(yaml_contents, Loader=yaml.Loader) def run_and_load_example_problem_yaml(cmd): subprocess.run(cmd, check=True) yaml_contents = example_problem.read_text() return yaml.load(yaml_contents, Loader=yaml.Loader) # fails with pytest.raises(subprocess.CalledProcessError): # overspecify tile size is bad subprocess.run( [ gemm, "example", example, "--arch=gfx950", "--wgts=128x128x128", "--mac_M=256", ], check=True, ) with pytest.raises(subprocess.CalledProcessError): subprocess.run( [ gemm, "example", example, "--arch=gfx950", "--wgs=128x2", "--workgroup_size_x=256", ], check=True, ) # PreSwizzleScaleGFX950 requires pretileScale; client must assert and exit non-zero with pytest.raises(subprocess.CalledProcessError): subprocess.run( [ gemm, "example", example, "--arch=gfx950", "--scaleSkipPermlane=PreSwizzleScaleGFX950", "--scale_A=Separate", "--scale_B=Separate", "--scaleBlockSize=32", "--sts=32x8/32x8", ], check=True, ) # PreSwizzleScaleGFX950 requires swizzleTileSize (m=32, n=32, k=8); wrong sts must fail with pytest.raises(subprocess.CalledProcessError): subprocess.run( [ gemm, "example", example, "--arch=gfx950", "--scaleSkipPermlane=PreSwizzleScaleGFX950", "--scale_A=Separate", "--scale_B=Separate", "--scaleBlockSize=32", "--sts=64x8/64x8", "--pretileScale", ], check=True, ) # setting tile size via shortcut post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--wgts=1024x2048x4096"] ) assert post["mac_m"] == 1024 assert post["mac_n"] == 2048 assert post["mac_k"] == 4096 # setting wg size via shortcut post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--wgs=64x4"] ) assert post["workgroup_size_x"] == 64 assert post["workgroup_size_y"] == 4 # setting mi via shortcut post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--mi=2x4x8"] ) assert post["wave_m"] == 2 assert post["wave_n"] == 4 assert post["wave_k"] == 8 assert post["wave_b"] == 1 post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--mi=4x8x16x2"] ) assert post["wave_m"] == 4 assert post["wave_n"] == 8 assert post["wave_k"] == 16 assert post["wave_b"] == 2 # setting lds options post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--lds=AB"] ) assert post["load_A"] == "BufferToLDSViaVGPR" assert post["load_B"] == "BufferToLDSViaVGPR" assert post["store"] == "VGPRToGlobalMemoryWithBuffer" post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--lds=BD"] ) assert post["load_A"] == "BufferToVGPR" assert post["load_B"] == "BufferToLDSViaVGPR" assert post["store"] == "VGPRToGlobalMemoryViaLDSWithBuffer" # setting d2l options post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--d2lds=AB"] ) assert post["load_A"] == "BufferToLDS" assert post["load_B"] == "BufferToLDS" post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--d2lds=A"] ) assert post["load_A"] == "BufferToLDS" assert post["load_B"] == "BufferToVGPR" post = run_and_load_example_yaml( [ gemm, "example", example, "--arch=gfx950", "--padLDS_A=22,33", "--padLDS_B=44,55", ] ) assert post["padLDS_A"] == [22, 33] assert post["padLDS_B"] == [44, 55] # setting mxlds options post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--mxlds=AB"] ) assert post["loadScale_A"] == "BufferToLDSViaVGPR" assert post["loadScale_B"] == "BufferToLDSViaVGPR" post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--mxlds=B"] ) assert post["loadScale_A"] == "BufferToVGPR" assert post["loadScale_B"] == "BufferToLDSViaVGPR" post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--mxd2lds=AB"] ) assert post["loadScale_A"] == "BufferToLDS" assert post["loadScale_B"] == "BufferToLDS" # setting swizzle tile size post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--sts=5x7/11x13"] ) assert post["swizzleTileSize"]["m"] == 5 assert post["swizzleTileSize"]["k"] == 7 assert post["swizzleTileSize"]["n"] == 11 assert post["swizzleTileSize"]["l"] == 13 # can also use a big X post = run_and_load_example_yaml( [gemm, "example", example, "--arch=gfx950", "--sts=5x7X11x13"] ) assert post["swizzleTileSize"]["m"] == 5 assert post["swizzleTileSize"]["k"] == 7 assert post["swizzleTileSize"]["n"] == 11 assert post["swizzleTileSize"]["l"] == 13 # pretileA post = run_and_load_example_yaml( [gemm, "example", example, "--wgts=256x256x256", "--pretileA=64x128"] ) assert post["types"]["pretileA"] == [64, 128] # pretileB post = run_and_load_example_yaml( [gemm, "example", example, "--wgts=256x256x256", "--pretileB=64x128"] ) assert post["types"]["pretileB"] == [64, 128] # setting data initialization modes post = run_and_load_example_problem_yaml( [ gemm, "exampleProblem", example_problem, "--arch=gfx950", "--initMode_A=Bounded", "--initMode_B=BoundedAlternatingSign", "--initMode_C=Unbounded", ] ) assert post["initMode_A"] == "DataInitMode(Bounded)" assert post["initMode_B"] == "DataInitMode(BoundedAlternatingSign)" assert post["initMode_C"] == "DataInitMode(Unbounded)" post = run_and_load_example_problem_yaml( [ gemm, "exampleProblem", example_problem, "--arch=gfx950", "--initMode_A=Identity", "--initMode_B=Ones", "--initMode_C=Zeros", ] ) assert post["initMode_A"] == "DataInitMode(Identity)" assert post["initMode_B"] == "DataInitMode(Ones)" assert post["initMode_C"] == "DataInitMode(Zeros)" mean_B = 0.0 std_dev_B = 1.0 mean_C = 2.0 std_dev_C = 3.0 post = run_and_load_example_problem_yaml( [ gemm, "exampleProblem", example_problem, "--arch=gfx950", "--initMode_A=TrigonometricFromFloat", f"--initMode_B=NormalFromFloat({mean_B}, {std_dev_B})", f"--initMode_C=NormalFromFloat({mean_C}, {std_dev_C})", ] ) assert post["initMode_A"] == "DataInitMode(TrigonometricFromFloat)" initMode_B = post["initMode_B"] assert initMode_B.startswith("DataInitMode(NormalFromFloat(") mean, std_dev = initMode_B.split("(")[-1][:-2].split(", ") assert float(mean) == mean_B assert float(std_dev) == std_dev_B initMode_C = post["initMode_C"] assert initMode_C.startswith("DataInitMode(NormalFromFloat(") mean, std_dev = initMode_C.split("(")[-1][:-2].split(", ") assert float(mean) == mean_C assert float(std_dev) == std_dev_C def test_gemm_generate_from_example(tmp_path): """GEMM 'generate' from the output of 'example'.""" example = tmp_path / "example.yaml" subprocess.run([gemm, "example", example], check=True) # We should be able to generate a kernel from the config file subprocess.run([gemm, "generate", "--config", example], check=True) def test_gemm_config(tmp_path): """GEMM load from config file.""" solution_params = ["--arch", "gfx90a", "--config", DP_HGEMM] solution_params = write_solution_config_if_present(tmp_path, solution_params) co_path = tmp_path / "test_config.co" yaml_path = co_path.with_suffix(".yaml") cmd = [gemm] cmd.extend(["generate", "--co", co_path]) cmd.extend(solution_params) subprocess.run(cmd, check=True) yaml_contents = yaml_path.read_text() client = yaml.load(yaml_contents, Loader=yaml.Loader) del client["version"] reference = yaml.load(DP_HGEMM, Loader=yaml.Loader) assert client == reference def test_gemm_generate(tmp_path): """GEMM 'generate' basics.""" with chdir(tmp_path): # "gemm generate" should pass subprocess.run([gemm, "generate"], check=True) # "gemm generate --asm" should write an assembly and two yaml files in the current directory before = list(tmp_path.glob("*.s")) + list(tmp_path.glob("*.yaml")) subprocess.run([gemm, "generate", "--asm"], check=True) after = list(tmp_path.glob("*.s")) + list(tmp_path.glob("*.yaml")) assert len(after) == len(before) + 3 # "gemm generate --asm test.s" should write an .s and .yaml pair asm_path = tmp_path / "test_asm.s" yaml_path = asm_path.with_suffix(".yaml") subprocess.run([gemm, "generate", "--asm", asm_path], check=True) assert asm_path.exists() assert yaml_path.exists() # possible to not write a pair? # "gemm generate --co" should write an co and two yaml files in the current directory before = list(tmp_path.glob("*.co")) + list(tmp_path.glob("*.yaml")) subprocess.run([gemm, "generate", "--co"], check=True) after = list(tmp_path.glob("*.co")) + list(tmp_path.glob("*.yaml")) assert len(after) == len(before) + 3 # "gemm generate --co test.co" should write .co+.yaml pair co_path = tmp_path / "test_co.co" yaml_path = asm_path.with_suffix(".yaml") subprocess.run([gemm, "generate", "--co", co_path], check=True) assert co_path.exists() assert yaml_path.exists() # "gemm generate --config" should fail with pytest.raises(subprocess.CalledProcessError): subprocess.run([gemm, "generate", "--config"], check=True) def test_gemm_validate(tmp_path): """GEMM generate and validate using one and two stages. This runs each problem/solution three times. """ isGFX120X = rocm_gfx().startswith("gfx120") problem_params = [["--m", "512", "--n", "512", "--k", "256", "--numWGs", "4"]] solution_params = [ # data-parallel gemm, float, params from command line # [], # data-parallel gemm, float, params from config file ["--config", DP_HGEMM_GFX120X if isGFX120X else DP_GEMM], # streamk gemm, float, params from command line # ["--streamk"], ] for problem, solution in itertools.product(problem_params, solution_params): gemm_validate_single_stage(tmp_path, solution, problem) gemm_validate_two_stage_codeobject(tmp_path, solution, problem) gemm_validate_two_stage_assembly(tmp_path, solution, problem) @pytest.mark.parametrize( "solution_params,problem_params", itertools.product(build_solution_params(), build_problem_params()), ) def test_gemm_validate_once(tmp_path, solution_params, problem_params): """GEMM generate (always) and validate (if arch matches).""" gemm_validate_two_stage_codeobject(tmp_path, solution_params, problem_params) @pytest.mark.parametrize("solution_params,problem_params", build_wgm_params()) def test_gemm_wgm(tmp_path, solution_params, problem_params): # TODO This is a temporary fix to enable GFX12 CI if rocm_gfx().startswith("gfx12"): return gemm_validate_single_stage(tmp_path, solution_params, problem_params) def test_kernel_graph_dot_truncation(tmp_path): """Validate Graphviz DOT rendering succeeds when node labels are truncated. - With truncation enabled (small max label length), kgraph.py should succeed and produce non-empty outputs. - With truncation disabled (0), kgraph.py should report a parse error. """ arch = rocm_gfx() if arch is not None and arch.startswith("gfx12"): pytest.skip("Skipping KernelGraph DOT truncation test on gfx12") if not gemm.exists(): pytest.skip("rocroller-gemm binary not found") kgraph = (Path(__file__).parent.parent / "scripts" / "kgraph.py").resolve() if not kgraph.exists(): pytest.skip("kgraph.py script not found") if shutil.which("dot") is None: pytest.skip("Graphviz 'dot' not available in PATH") def run_cmd(cmd, env=None, cwd=None): return subprocess.run(cmd, cwd=cwd, env=env, text=True, capture_output=True) def assert_non_empty(path: Path): assert path.exists(), f"Expected file to exist: {path}" assert path.stat().st_size > 0, f"Expected file to be non-empty: {path}" def generate_asm(asm_path: Path, extra_env: dict): env = os.environ.copy() env.update(extra_env) cmd = [ str(gemm), "--workgroupMappingDim=0", "--workgroupMappingValue=6", "generate", "--asm", str(asm_path), ] p = run_cmd(cmd, env=env, cwd=tmp_path) if p.returncode == SOLUTION_NOT_SUPPORTED_ON_ARCH: pytest.skip("GEMM solution not supported on this architecture") assert p.returncode == 0, ( "rocroller-gemm failed\n" f"cmd: {cmd}\n" f"stdout:\n{p.stdout}\n" f"stderr:\n{p.stderr}\n" ) assert_non_empty(asm_path) def run_kgraph(asm_path: Path, pdf_path: Path): cmd = [str(kgraph), str(asm_path), "-o", str(pdf_path)] p = run_cmd(cmd, env=os.environ.copy(), cwd=tmp_path) combined = (p.stdout or "") + "\n" + (p.stderr or "") return p, combined # Case 1 : truncation enabled(should succeed) asm_trunc = tmp_path / "workgroupmapping_truncated5.s" pdf_trunc = tmp_path / "workgroupmapping_truncated5.pdf" generate_asm( asm_trunc, { "ROCROLLER_SERIALIZE_KERNEL_GRAPH_DOT": "1", "ROCROLLER_KGRAPH_NODE_LABEL_MAX_LENGTH": "5", }, ) p, combined = run_kgraph(asm_trunc, pdf_trunc) assert p.returncode == 0, ( "kgraph.py expected to succeed with truncation enabled\n" f"output:\n{combined}\n" ) dot_trunc = pdf_trunc.with_suffix(".dot") norm_trunc = pdf_trunc.with_name(pdf_trunc.stem + "_normalized.s") assert_non_empty(dot_trunc) assert_non_empty(pdf_trunc) assert_non_empty(norm_trunc) # Case 2 : truncation disabled(should error in kgraph parse) asm_untrunc = tmp_path / "workgroupmapping_untruncated.s" pdf_untrunc = tmp_path / "workgroupmapping_untruncated.pdf" generate_asm( asm_untrunc, { "ROCROLLER_SERIALIZE_KERNEL_GRAPH_DOT": "1", "ROCROLLER_KGRAPH_NODE_LABEL_MAX_LENGTH": "0", }, ) p2, combined2 = run_kgraph(asm_untrunc, pdf_untrunc) assert ( p2.returncode != 0 or "syntax error" in combined2 or "longer than 16384" in combined2 ), ( "Expected Graphviz parse error when truncation is disabled\n" f"returncode: {p2.returncode}\n" f"output:\n{combined2}\n" ) dot_untrunc = pdf_untrunc.with_suffix(".dot") assert_non_empty(dot_untrunc) max_line2 = max( len(line) for line in dot_untrunc.read_text(errors="ignore").splitlines() or [""] ) assert ( max_line2 >= 16384 ), f"Expected an extremely long DOT line without truncation, got max {max_line2}" if __name__ == "__main__": print("Solution params") for p in build_solution_params(): print(p) print("Problem params") for p in build_problem_params(): print(p)