[MLIR] Rework generate-test-checks.py to attach CHECK lines to the source (test) file.

[timshen]

Summary

This patch adds --source flag to indicate the source file. Then it tries to find insert
points in the source file and insert corresponding checks at those places.

Example output from Tensorflow XLA:

// -----

// CHECK-LABEL:   func @main.3(
// CHECK-SAME:                 %[[VAL_0:.*]]: memref<2x2xf32> {xla_lhlo.params = 0 : index},
// CHECK-SAME:                 %[[VAL_1:.*]]: memref<16xi8> {xla_lhlo.alloc = 0 : index, xla_lhlo.liveout = true}) {
// CHECK:           %[[VAL_2:.*]] = constant 0 : index
// CHECK:           %[[VAL_3:.*]] = constant 0 : index
// CHECK:           %[[VAL_4:.*]] = std.view %[[VAL_1]]{{\[}}%[[VAL_3]]][] : memref<16xi8> to memref<2x2xf32>
// CHECK:           "xla_lhlo.tanh"(%[[VAL_0]], %[[VAL_4]]) : (memref<2x2xf32>, memref<2x2xf32>) -> ()
// CHECK:           return
// CHECK:         }
func @main(%value0: tensor<2x2xf32>) -> tensor<2x2xf32> {
  %res = "xla_hlo.tanh"(%value0) : (tensor<2x2xf32>) -> tensor<2x2xf32>
  return %res : tensor<2x2xf32>
}

[timshen]

Fix typo.

[timshen]

Remove extra empty line in the end.

[mehdi_amini]

Nice improvement, thanks!

[rriddle]

Looks good, thanks!

mlir/utils/generate-test-checks.py

11–14

Please add documentation for any new behavior here.

59

nit: Add a function comment for these.

96

nit: Add comments.

[timshen]

Added documentation in succeeding commits.

mlir/utils/generate-test-checks.py

 #!/usr/bin/env python3
 """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
 stripped down variant.
 
 Example usage:
 $ generate-test-checks.py foo.mlir
 $ mlir-opt foo.mlir -transformation | generate-test-checks.py
 
 The script will heuristically insert 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. 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!
 """
 
 # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
class SSAVariableNamer:
   # Push a new variable name scope.
   def push_name_scope(self):
     self.scopes.append({})
 
   # Pop the last variable name scope.
   def pop_name_scope(self):
     self.scopes.pop()
 
+  def num_scopes(self):
+    return len(self.scopes)
+
+  def clear_counter(self):
+    self.name_counter = 0
+
 
 # Process a line of input that has been split at each SSA identifier '%'.
 def process_line(line_chunks, variable_namer):
   output_line = ''
 
   # Process the rest that contained an SSA value name.
   for chunk in line_chunks:
     m = SSA_RE.match(chunk)
else:
       output_line += '%[[' + variable + ':.*]]'
 
     # Append the non named group.
     output_line += chunk[len(ssa_name):]
 
   return output_line + '\n'
 
 
+def process_source_lines(source_lines, note, args):
+  source_split_re = re.compile(args.source_delim_regex)
+
+  source_segments = [[]]
+  for line in source_lines:
+    if line == note:
+      continue
+    if line.find(args.check_prefix) != -1:
+      continue
+    if source_split_re.search(line):
+      source_segments.append([])
+
+    source_segments[-1].append(line + '\n')
+  return source_segments
+
+
 # Pre-process a line of input to remove any character sequences that will be
 # problematic with FileCheck.
 def preprocess_line(line):
   # Replace any double brackets, '[[' with escaped replacements. '[['
   # corresponds to variable names in FileCheck.
   output_line = line.replace('[[', '{{\\[\\[}}')
 
   # Replace any single brackets that are followed by an SSA identifier, the
parser = argparse.ArgumentParser(
       description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
   parser.add_argument(
       '--check-prefix', default='CHECK', help='Prefix to use from check file.')
   parser.add_argument(
       '-o',
       '--output',
       nargs='?',
       type=argparse.FileType('w'),
-      default=sys.stdout)
+      default=None)
   parser.add_argument(
       'input',
       nargs='?',
       type=argparse.FileType('r'),
       default=sys.stdin)
+  parser.add_argument(
+      '--source', type=str,
+      help='Print each CHECK chunk before each delimeter line in the source'
+           'file, respectively. The delimeter lines are identified by '
+           '--source_delim_regex.')
+  parser.add_argument('--source_delim_regex', type=str, default='func @')
+  parser.add_argument(
+      '--starts_from_scope', type=int, default=1,
+      help='Omit the top specified level of content. For example, by default '
+           'it omits "module {"')
+  parser.add_argument('-i', '--inplace', action='store_true', default=False)
+
   args = parser.parse_args()
 
   # Open the given input file.
   input_lines = [l.rstrip() for l in args.input]
   args.input.close()
 
-  output_lines = []
-
   # Generate a note used for the generated check file.
   script_name = os.path.basename(__file__)
   autogenerated_note = (ADVERT + 'utils/' + script_name)
-  output_lines.append(autogenerated_note + '\n')
 
+  source_segments = None
+  if args.source:
+    source_segments = process_source_lines(
+        [l.rstrip() for l in open(args.source, 'r')],
+        autogenerated_note,
+        args
+    )
+
+  if args.inplace:
+    assert args.output is None
+    output = open(args.source, 'w')
+  elif args.output is None:
+    output = sys.stdout
+  else:
+    output = args.output
+
+  output_segments = [[]]
   # A map containing data used for naming SSA value names.
   variable_namer = SSAVariableNamer()
   for input_line in input_lines:
     if not input_line:
       continue
     lstripped_input_line = input_line.lstrip()
 
     # Lines with blocks begin with a ^. These lines have a trailing comment
     # that needs to be stripped.
     is_block = lstripped_input_line[0] == '^'
     if is_block:
       input_line = input_line.rsplit('//', 1)[0].rstrip()
 
-    # Top-level operations are heuristically the operations at nesting level 1.
-    is_toplevel_op = (not is_block and input_line.startswith('  ') and
-                      input_line[2] != ' ' and input_line[2] != '}')
+    cur_level = variable_namer.num_scopes()
 
     # If the line starts with a '}', pop the last name scope.
     if lstripped_input_line[0] == '}':
       variable_namer.pop_name_scope()
+      cur_level = variable_namer.num_scopes()
 
     # If the line ends with a '{', push a new name scope.
     if input_line[-1] == '{':
       variable_namer.push_name_scope()
+      if cur_level == args.starts_from_scope:
+        output_segments.append([])
+
+    # Omit lines at the near top level e.g. "module {".
+    if cur_level < args.starts_from_scope:
+      continue
+
+    if len(output_segments[-1]) == 0:
+      variable_namer.clear_counter()
 
     # Preprocess the input to remove any sequences that may be problematic with
     # FileCheck.
     input_line = preprocess_line(input_line)
 
     # Split the line at the each SSA value name.
     ssa_split = input_line.split('%')
 
     # If this is a top-level operation use 'CHECK-LABEL', otherwise 'CHECK:'.
-    if not is_toplevel_op or not ssa_split[0]:
+    if len(output_segments[-1]) != 0 or not ssa_split[0]:
       output_line = '// ' + args.check_prefix + ': '
       # Pad to align with the 'LABEL' statements.
       output_line += (' ' * len('-LABEL'))
 
       # Output the first line chunk that does not contain an SSA name.
       output_line += ssa_split[0]
 
       # Process the rest of the input line.
       output_line += process_line(ssa_split[1:], variable_namer)
 
     else:
-      # Append a newline to the output to separate the logical blocks.
-      output_lines.append('\n')
-      output_line = '// ' + args.check_prefix + '-LABEL: '
-
       # Output the first line chunk that does not contain an SSA name for the
       # label.
-      output_line += ssa_split[0] + '\n'
+      output_line = '// ' + args.check_prefix + '-LABEL: ' + ssa_split[0] + '\n'
 
-      # Process the rest of the input line on a separate check line.
-      if len(ssa_split) > 1:
+      # Process the rest of the input line on separate check lines.
+      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])
 
         # Process the rest of the line.
-        output_line += process_line(ssa_split[1:], variable_namer)
+        output_line += process_line([argument], variable_namer)
 
     # Append the output line.
-    output_lines.append(output_line)
+    output_segments[-1].append(output_line)
+
+  output.write(autogenerated_note + '\n')
 
   # Write the output.
-  for output_line in output_lines:
-    args.output.write(output_line)
-  args.output.write('\n')
-  args.output.close()
+  if source_segments:
+    assert len(output_segments) == len(source_segments)
+    for check_segment, source_segment in zip(output_segments, source_segments):
+      for line in check_segment:
+        output.write(line)
+      for line in source_segment:
+        output.write(line)
+  else:
+    for segment in output_segments:
+      output.write('\n')
+      for output_line in segment:
+        output.write(output_line)
+    output.write('\n')
+  output.close()
 
 
 if __name__ == '__main__':
   main()