Add --verbose flag to verify_codegen.py

Author: stijnh

.github/workflows/cmake-run-cuda.yml

@@ -34,6 +34,12 @@ jobs:
       with:
         submodules: 'true'
 
+    # Run codegen tests first, so things fail immediately before building the full project
+    - name: Codegen tests
+      # These check exact instruction counts in the generated PTX, which can legitimately
+      # differ between nvcc versions even when the code itself is correct.
+      run: python3 tests_codegen/verify_codegen.py tests_codegen/*.cu --nvcc nvcc
+
     - name: Configure CMake
       # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make.
       # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type
@@ -48,8 +54,3 @@ jobs:
       # Execute tests defined by the CMake configuration.
       # See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail
       run: ./tests/kernel_float_tests --durations=yes --success --verbosity=high ~[GPU]
-
-    - name: Codegen instruction checks
-      # These check exact instruction counts in the generated PTX, which can legitimately
-      # differ between nvcc versions even when the code itself is correct.
-      run: python3 tests_codegen/verify_codegen.py tests_codegen/*.cu --nvcc nvcc

tests_codegen/bf16_ops.cu

@@ -81,9 +81,8 @@ __global__ void bfloat16_eq(const kf::bfloat16x4 *input, bool *output) {
     output[0] = kf::all(*input == *input);
 }
 
-// Division is not packed at all: each of the 4 lanes is extracted individually and divided
-// with a scalar `div.rn.f32`, then converted back with `cvt.rn.bf16.f32`.
-// CHECK-COUNT-4: div.rn.f32
+// Division is not packed at all: each of the 4 lanes is extracted individually.
+// CHECK-COUNT-4: div.(rn|approx).f32
 // CHECK-NOT: div.bf16x2
 __global__ void bfloat16_div(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
     *output = *input / *input;

tests_codegen/verify_codegen.py

@@ -54,38 +54,33 @@
 
 
 class Kernel:
-    def __init__(self, name):
+    def __init__(self, name, requires):
         self.name = name
-        self.requires = []  # list of (pattern, expected_count)
-        self.forbids = []
+        self.requires = requires  # list of (pattern, expected_count)
 
 
 def parse_kernels(source: Path):
     kernels = []
-    pending_requires, pending_forbids = [], []
+    pending_requires = []
     arch = None
 
     for line in source.read_text().splitlines():
         if m := CHECK_COUNT_RE.search(line):
             pending_requires.append((m.group(2).strip(), int(m.group(1))))
         elif m := CHECK_RE.search(line):
-            # `None` means "at least once" (exact count is not checked)
             pending_requires.append((m.group(1).strip(), 1))
         elif m := CHECK_NOT_RE.search(line):
-            pending_forbids.append(m.group(1).strip())
+            pending_requires.append((m.group(1).strip(), 0))
         elif m := ARCH_RE.search(line):
             arch = m.group(1)
         elif m := KERNEL_RE.search(line):
-            kernel = Kernel(m.group(1))
-            kernel.requires = pending_requires
-            kernel.forbids = pending_forbids
-            kernels.append(kernel)
-            pending_requires, pending_forbids = [], []
+            kernels.append(Kernel(m.group(1), pending_requires))
+            pending_requires = []
 
     return kernels, arch
 
 
-def get_ptx(nvcc, source, arch):
+def get_ptx(nvcc, source, arch, verbose=False):
     kf_root = Path(__file__).parent.parent.resolve() / "include"
     cmd = [
         nvcc,
@@ -99,6 +94,9 @@ def get_ptx(nvcc, source, arch):
         str(source),
     ]
 
+    if verbose:
+        print(f"$ {' '.join(cmd)}", file=sys.stderr)
+
     result = subprocess.run(cmd, capture_output=True, text=True)
     if result.returncode != 0:
         print(f"command failed: {' '.join(cmd)}", file=sys.stderr)
@@ -128,25 +126,49 @@ def extract_ptx_block(ptx: str, kernel_name: str) -> str:
     return None
 
 
-def check_file(source: Path, nvcc: str, arch: str) -> "tuple[int, int]":
-    """Returns (num_kernels, num_failed) for this source file."""
-    kernels, arch_directive = parse_kernels(source)
+def check_file(
+    source: Path,
+    nvcc: str,
+    arch: str,
+    verbose: bool = False,
+    kernel_filter: "set[str] | None" = None,
+) -> "tuple[int, int, int]":
+    """Returns (num_passed, num_failed, num_skipped) for this source file."""
+    all_kernels, arch_directive = parse_kernels(source)
     arch = arch_directive or arch
 
-    kernels = [k for k in kernels if k.requires or k.forbids]
-    if not kernels:
+    checkable = [k for k in all_kernels if k.requires]
+    if not checkable:
         print(
             f"error: {source} has no kernel preceded by "
             "// CHECK:, // CHECK-COUNT-<N>:, or // CHECK-NOT: directives",
             file=sys.stderr,
         )
-        return 0, 1
+        return 0, 1, 0
+
+    # Kernels with no directives at all are skipped silently; they were never meant to be
+    # checked. Kernels excluded by --kernel are also skipped, but reported below.
+    num_skipped = len(all_kernels) - len(checkable)
+
+    if kernel_filter is not None:
+        kernels = [k for k in checkable if k.name in kernel_filter]
+        num_skipped += len(checkable) - len(kernels)
+        for k in checkable:
+            if k.name not in kernel_filter:
+                print(f"SKIP: {k.name} (excluded by --kernel)")
+    else:
+        kernels = checkable
+
+    if not kernels:
+        # Nothing in this file matches --kernel; not an error, just nothing to do here.
+        return 0, 0, num_skipped
 
-    ptx = get_ptx(nvcc, source, arch)
+    ptx = get_ptx(nvcc, source, arch, verbose=verbose)
     if ptx is None:
         print(f"FAIL: {source}: nvcc failed to compile this file", file=sys.stderr)
-        return len(kernels), len(kernels)
+        return 0, len(kernels), num_skipped
 
+    num_passed = 0
     num_failed = 0
 
     for kernel in kernels:
@@ -162,18 +184,24 @@ def check_file(source: Path, nvcc: str, arch: str) -> "tuple[int, int]":
 
         failures = []
         for pattern, expected_count in kernel.requires:
-            if expected_count is None:
-                if not re.search(pattern, block, re.MULTILINE):
-                    failures.append(f"CHECK pattern not found: {pattern!r}")
-            else:
-                actual_count = len(re.findall(pattern, block, re.MULTILINE))
-                if actual_count != expected_count:
-                    failures.append(
-                        f"CHECK pattern found {actual_count} time(s), expected {expected_count}: {pattern!r}"
-                    )
-        for pattern in kernel.forbids:
-            if re.search(pattern, block, re.MULTILINE):
-                failures.append(f"CHECK-NOT pattern found: {pattern!r}")
+            actual_count = len(re.findall(pattern, block, re.MULTILINE))
+            ok = (
+                actual_count >= 1
+                if expected_count is None
+                else actual_count == expected_count
+            )
+            want = "at least 1" if expected_count is None else str(expected_count)
+
+            if verbose:
+                status = "success" if ok else "FAILURE"
+                print(
+                    f"[{kernel.name}] {status} (want {want}, found {actual_count}): {pattern!r}",
+                    file=sys.stderr,
+                )
+            if not ok:
+                failures.append(
+                    f"pattern found {actual_count} time(s), expected {want}: {pattern!r}"
+                )
 
         if failures:
             num_failed += 1
@@ -184,9 +212,16 @@ def check_file(source: Path, nvcc: str, arch: str) -> "tuple[int, int]":
             )
             print(block, file=sys.stderr)
         else:
+            num_passed += 1
             print(f"OK: {kernel.name} (ptx, {arch})")
+            if verbose:
+                print(
+                    f"----- generated ptx for {kernel.name} ({arch}) -----",
+                    file=sys.stderr,
+                )
+                print(block, file=sys.stderr)
 
-    return len(kernels), num_failed
+    return num_passed, num_failed, num_skipped
 
 
 def main():
@@ -195,22 +230,52 @@ def main():
     )
     parser.add_argument("source", type=Path, nargs="+")
     parser.add_argument("--nvcc", default="nvcc")
-    parser.add_argument("--verbose", action="store_true")
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="print the nvcc command, per-pattern match results, and the generated PTX "
+        "for every kernel, not just the ones that fail",
+    )
     parser.add_argument(
         "--arch", help="used unless overridden by a // ARCH: directive", default="sm_80"
     )
+    parser.add_argument(
+        "--kernel",
+        action="append",
+        help="only check the kernel with this name, skipping all others (repeatable)",
+    )
     args = parser.parse_args()
 
-    total_kernels = 0
+    kernel_filter = set(args.kernel) if args.kernel else None
+
+    total_passed = 0
     total_failed = 0
+    total_skipped = 0
 
     for source in args.source:
         print(f"=== {source} ===")
-        num_kernels, num_failed = check_file(source, args.nvcc, args.arch)
-        total_kernels += num_kernels
+        num_passed, num_failed, num_skipped = check_file(
+            source,
+            args.nvcc,
+            args.arch,
+            verbose=args.verbose,
+            kernel_filter=kernel_filter,
+        )
+        total_passed += num_passed
         total_failed += num_failed
+        total_skipped += num_skipped
 
-    print(f"TEST PASSED: {total_kernels - total_failed} / {total_kernels}")
+    if kernel_filter is not None and total_passed + total_failed == 0:
+        print(
+            f"error: no kernel named {sorted(kernel_filter)} found in any source file",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    total = total_passed + total_failed + total_skipped
+    print(
+        f"RESULT: {total_passed} passed, {total_failed} failed, {total_skipped} skipped ({total} total)"
+    )
 
     if total_failed:
         sys.exit(1)