[Mlir-commits] [mlir] [mlir][utils] Update generate-test-checks.py (PR #136757)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/136757

At the moment, the `CHECK-SAME` lines generated by
"generate-test-checks.py" (i.e. check-lines that correspond to the
preceeding `CHECK-LABEL` line) are indented to match the label length.
For example,

```mlir
func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(%arg0: memref<5xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<3x7xf32>) {
  linalg.batch_reduce_matmul indexing_maps = (...)
}
```

will lead to the following:

```mlir
// CHECK-LABEL:   func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(
// CHECK-SAME:                                                                  %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<5xf32>,
// CHECK-SAME:                                                                  %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<2x5x7xf32>,
// CHECK-SAME:                                                                  %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<3x7xf32>) {
// CHECK:           linalg.batch_reduce_matmul indexing_maps = (...)
```

This indentation is unnecasarilly deep. With this change, for labales
that are longer than 20 chars, the indentation is trimmed to 4 spaces:
```mlir
// CHECK-LABEL:   func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(
// CHECK-SAME:        %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<5xf32>,
// CHECK-SAME:        %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<2x5x7xf32>,
// CHECK-SAME:        %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<3x7xf32>) {
// CHECK:           linalg.batch_reduce_matmul indexing_maps = (...)
```


>From 8ccc12a8e2c631fee8eee54fcb2917eab8ed7d13 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 22 Apr 2025 20:06:16 +0000
Subject: [PATCH] [mlir][utils] Update generate-test-checks.py

At the moment, the `CHECK-SAME` lines generated by
"generate-test-checks.py" (i.e. check-lines that correspond to the
preceeding `CHECK-LABEL` line) are indented to match the label length.
For example,

```mlir
func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(%arg0: memref<5xf32>, %arg1: memref<2x5x7xf32>, %arg2: memref<3x7xf32>) {
  linalg.batch_reduce_matmul indexing_maps = (...)
}
```

will lead to the following:

```mlir
// CHECK-LABEL:   func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(
// CHECK-SAME:                                                                  %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<5xf32>,
// CHECK-SAME:                                                                  %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<2x5x7xf32>,
// CHECK-SAME:                                                                  %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<3x7xf32>) {
// CHECK:           linalg.batch_reduce_matmul indexing_maps = (...)
```

This indentation is unnecasarilly deep. With this change, for labales
that are longer than 20 chars, the indentation is trimmed to 4 spaces:
```mlir
// CHECK-LABEL:   func.func @batch_reduce_matmul_bcast_k_to_fill_missing_dims_A(
// CHECK-SAME:        %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<5xf32>,
// CHECK-SAME:        %[[VAL_1:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<2x5x7xf32>,
// CHECK-SAME:        %[[VAL_2:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: memref<3x7xf32>) {
// CHECK:           linalg.batch_reduce_matmul indexing_maps = (...)
```
---
 mlir/utils/generate-test-checks.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mlir/utils/generate-test-checks.py b/mlir/utils/generate-test-checks.py
index f1dd7a2843893..d157af9c3cab7 100755
--- a/mlir/utils/generate-test-checks.py
+++ b/mlir/utils/generate-test-checks.py
@@ -408,8 +408,12 @@ def main():
             for argument in ssa_split[1:]:
                 output_line += "// " + args.check_prefix + "-SAME:  "
 
-                # Pad to align with the original position in the line.
-                output_line += " " * len(ssa_split[0])
+                # Pad to align with the original position in the line (i.e. where the label ends),
+                # unless the label is more than 20 chars long, in which case pad with 4 spaces
+                # (this is to avoid deep indentation).
+                label_length = len(ssa_split[0])
+                pad_depth = label_length if label_length < 21 else 4
+                output_line += " " * pad_depth
 
                 # Process the rest of the line.
                 output_line += process_line(