xref: /llvm-project/mlir/utils/generate-test-checks.py (revision f28c5aca1766889cc19b09ee32796df3e7b52cfe) 
1#!/usr/bin/env python
2"""A script to generate FileCheck statements for mlir unit tests.
3
4This script is a utility to add FileCheck patterns to an mlir file.
5
6NOTE: The input .mlir is expected to be the output from the parser, not a
7stripped down variant.
8
9Example usage:
10$ generate-test-checks.py foo.mlir
11$ mlir-opt foo.mlir -transformation | generate-test-checks.py
12
13The script will heuristically insert CHECK/CHECK-LABEL commands for each line
14within the file. By default this script will also try to insert string
15substitution blocks for all SSA value names. The script is designed to make
16adding checks to a test case fast, it is *not* designed to be authoritative
17about what constitutes a good test!
18"""
19
20# Copyright 2019 The MLIR Authors.
21#
22# Licensed under the Apache License, Version 2.0 (the "License");
23# you may not use this file except in compliance with the License.
24# You may obtain a copy of the License at
25#
26#   http://www.apache.org/licenses/LICENSE-2.0
27#
28# Unless required by applicable law or agreed to in writing, software
29# distributed under the License is distributed on an "AS IS" BASIS,
30# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
31# See the License for the specific language governing permissions and
32# limitations under the License.
33
34import argparse
35import os  # Used to advertise this file's name ("autogenerated_note").
36import re
37import sys
38
39ADVERT = '// NOTE: Assertions have been autogenerated by '
40
41# Regex command to match an SSA identifier.
42SSA_RE_STR = '[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*'
43SSA_RE = re.compile(SSA_RE_STR)
44
45
46# Class used to generate and manage string substitution blocks for SSA value
47# names.
48class SSAVariableNamer:
49
50  def __init__(self):
51    self.scopes = []
52    self.name_counter = 0
53
54  # Generate a substitution name for the given ssa value name.
55  def generate_name(self, ssa_name):
56    variable = 'VAL_' + str(self.name_counter)
57    self.name_counter += 1
58    self.scopes[-1][ssa_name] = variable
59    return variable
60
61  # Push a new variable name scope.
62  def push_name_scope(self):
63    self.scopes.append({})
64
65  # Pop the last variable name scope.
66  def pop_name_scope(self):
67    self.scopes.pop()
68
69
70# Process a line of input that has been split at each SSA identifier '%'.
71def process_line(line_chunks, variable_namer):
72  output_line = ''
73
74  # Process the rest that contained an SSA value name.
75  for chunk in line_chunks:
76    m = SSA_RE.match(chunk)
77    ssa_name = m.group(0)
78
79    # Check if an existing variable exists for this name.
80    variable = None
81    for scope in variable_namer.scopes:
82      variable = scope.get(ssa_name)
83      if variable is not None:
84        break
85
86    # If one exists, then output the existing name.
87    if variable is not None:
88      output_line += '[[' + variable + ']]'
89    else:
90      # Otherwise, generate a new variable.
91      variable = variable_namer.generate_name(ssa_name)
92      output_line += '[[' + variable + ':%.*]]'
93
94    # Append the non named group.
95    output_line += chunk[len(ssa_name):]
96
97  return output_line + '\n'
98
99
100# Pre-process a line of input to remove any character sequences that will be
101# problematic with FileCheck.
102def preprocess_line(line):
103  # Replace any double brackets, '[[' with escaped replacements. '[['
104  # corresponds to variable names in FileCheck.
105  output_line = line.replace('[[', '{{\\[\\[}}')
106
107  # Replace any single brackets that are followed by an SSA identifier, the
108  # identifier will be replace by a variable; Creating the same situation as
109  # above.
110  output_line = output_line.replace('[%', '{{\\[}}%')
111
112  return output_line
113
114
115def main():
116  parser = argparse.ArgumentParser(
117      description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
118  parser.add_argument(
119      '--check-prefix', default='CHECK', help='Prefix to use from check file.')
120  parser.add_argument(
121      '-o',
122      '--output',
123      nargs='?',
124      type=argparse.FileType('w'),
125      default=sys.stdout)
126  parser.add_argument(
127      'input',
128      nargs='?',
129      type=argparse.FileType('r'),
130      default=sys.stdin)
131  args = parser.parse_args()
132
133  # Open the given input file.
134  input_lines = [l.rstrip() for l in args.input]
135  args.input.close()
136
137  output_lines = []
138
139  # Generate a note used for the generated check file.
140  script_name = os.path.basename(__file__)
141  autogenerated_note = (ADVERT + 'utils/' + script_name)
142  output_lines.append(autogenerated_note + '\n')
143
144  # A map containing data used for naming SSA value names.
145  variable_namer = SSAVariableNamer()
146  for input_line in input_lines:
147    if not input_line:
148      continue
149    lstripped_input_line = input_line.lstrip()
150
151    # Lines with blocks begin with a ^. These lines have a trailing comment
152    # that needs to be stripped.
153    is_block = lstripped_input_line[0] == '^'
154    if is_block:
155      input_line = input_line.rsplit('//', 1)[0].rstrip()
156
157    # Top-level operations are heuristically the operations at nesting level 1.
158    is_toplevel_op = (not is_block and input_line.startswith('  ') and
159                      input_line[2] != ' ' and input_line[2] != '}')
160
161    # If the line starts with a '}', pop the last name scope.
162    if lstripped_input_line[0] == '}':
163      variable_namer.pop_name_scope()
164
165    # If the line ends with a '{', push a new name scope.
166    if input_line[-1] == '{':
167      variable_namer.push_name_scope()
168
169    # Preprocess the input to remove any sequences that may be problematic with
170    # FileCheck.
171    input_line = preprocess_line(input_line)
172
173    # Split the line at the each SSA value name.
174    ssa_split = input_line.split('%')
175
176    # If this is a top-level operation use 'CHECK-LABEL', otherwise 'CHECK:'.
177    if not is_toplevel_op or not ssa_split[0]:
178      output_line = '// ' + args.check_prefix + ': '
179      # Pad to align with the 'LABEL' statements.
180      output_line += (' ' * len('-LABEL'))
181
182      # Output the first line chunk that does not contain an SSA name.
183      output_line += ssa_split[0]
184
185      # Process the rest of the input line.
186      output_line += process_line(ssa_split[1:], variable_namer)
187
188    else:
189      # Append a newline to the output to separate the logical blocks.
190      output_lines.append('\n')
191      output_line = '// ' + args.check_prefix + '-LABEL: '
192
193      # Output the first line chunk that does not contain an SSA name for the
194      # label.
195      output_line += ssa_split[0] + '\n'
196
197      # Process the rest of the input line on a separate check line.
198      if len(ssa_split) > 1:
199        output_line += '// ' + args.check_prefix + '-SAME:  '
200
201        # Pad to align with the original position in the line.
202        output_line += ' ' * len(ssa_split[0])
203
204        # Process the rest of the line.
205        output_line += process_line(ssa_split[1:], variable_namer)
206
207    # Append the output line.
208    output_lines.append(output_line)
209
210  # Write the output.
211  for output_line in output_lines:
212    args.output.write(output_line)
213  args.output.write('\n')
214  args.output.close()
215
216
217if __name__ == '__main__':
218  main()
219