[TOOLS] Improve generate-test-checks.py (#5300)

- Format the doc string using the `reStructuredText` format. 
- Lift the example instructions from the `.mlir.in` file to the
docstring. Previously we matched the `module` keyword twice and
encountered errors such as `assert len(output_segments) ==
len(source_segments),`. It's also fine to update the regex to something
like `\bmodule` to solve the problem, but I think lifting it from the
input file is just simpler.

Author: Jokeren

test/TritonGPU/samples/descriptor-matmul-pipeline.mlir

@@ -10,7 +10,11 @@
 // CHECK: #[[$ATTR_2:.+]] = #ttg.shared<{vec = 8, perPhase = 1, maxPhase = 8, order = [1, 0], hasLeadingOffset = true}>
 // CHECK: #[[$ATTR_3:.+]] = #ttg.shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], hasLeadingOffset = false}>
 // CHECK: #[[$ATTR_4:.+]] = #ttg.shared<{vec = 8, perPhase = 1, maxPhase = 8, order = [0, 1], hasLeadingOffset = true}>
-// RUN: triton-opt %s -split-input-file -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | FileCheck --dump-input-context=50 %s
+// To regenerate this test case, run the command
+//    triton-opt test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | \
+//    utils/generate-test-checks.py --source test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in --source_delim_regex="\bmodule" \
+//    -o test/TritonGPU/samples/descriptor-matmul-pipeline.mlir
+// RUN: triton-opt %s -split-input-file -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | FileCheck --dump-input-context=51 %s
 // CHECK-LABEL:   tt.func public @matmul_kernel_with_descriptors(
 // CHECK-SAME:  %[[VAL_0:.*]]: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %[[VAL_1:.*]]: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %[[VAL_2:.*]]: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %[[VAL_3:.*]]: i32 {tt.divisibility = 16 : i32}, %[[VAL_4:.*]]: i32 {tt.divisibility = 16 : i32}, %[[VAL_5:.*]]: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} {
 // CHECK:           %[[VAL_6:.*]] = arith.constant 2 : i32

test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in

@@ -1,8 +1,8 @@
 // To regenerate this test case, run the command
 //    triton-opt test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | \
-//    utils/generate-test-checks.py --source test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in --source_delim_regex="module" \
+//    utils/generate-test-checks.py --source test/TritonGPU/samples/descriptor-matmul-pipeline.mlir.in --source_delim_regex="\bmodule" \
 //    -o test/TritonGPU/samples/descriptor-matmul-pipeline.mlir
-// RUN: triton-opt %s -split-input-file -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | FileCheck --dump-input-context=50 %s
+// RUN: triton-opt %s -split-input-file -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | FileCheck --dump-input-context=51 %s
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} {
   tt.func public @matmul_kernel_with_descriptors(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} {
     %c8_i32 = arith.constant 8 : i32

utils/generate-test-checks.py

@@ -1,23 +1,29 @@
 #!/usr/bin/env python3
-"""A script to generate FileCheck statements for mlir unit tests.
+"""
+===============================================================
+A script to generate FileCheck statements for mlir unit tests.
+===============================================================
 
 This script is a utility to add FileCheck patterns to an mlir file.
 
-NOTE: The input .mlir is expected to be the output from the parser, not a
+NOTE: The input ``.mlir`` is expected to be the output from the parser, not a
 stripped down variant.
 
 Example usage:
-$ generate-test-checks.py foo.mlir
-$ mlir-opt foo.mlir -transformation | generate-test-checks.py
-$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir
-$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i
-$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i --source_delim_regex='gpu.func @'
+
+.. code-block:: shell
+
+    $ generate-test-checks.py foo.mlir
+    $ mlir-opt foo.mlir -transformation | generate-test-checks.py
+    $ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir
+    $ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i
+    $ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i --source_delim_regex='gpu.func @'
 
 The script will heuristically generate CHECK/CHECK-LABEL commands for each line
 within the file. By default this script will also try to insert string
-substitution blocks for all SSA value names. If --source file is specified, the
+substitution blocks for all SSA value names. If ``--source file`` is specified, the
 script will attempt to insert the generated CHECKs to the source file by looking
-for line positions matched by --source_delim_regex.
+for line positions matched by ``--source_delim_regex``.
 
 The script is designed to make adding checks to a test case fast, it is *not*
 designed to be authoritative about what constitutes a good test!