502 lines · python
1#!/usr/bin/env python32"""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 a7stripped down variant.8 9Example usage:10$ generate-test-checks.py foo.mlir11$ mlir-opt foo.mlir -transformation | generate-test-checks.py12$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir13$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i14$ mlir-opt foo.mlir -transformation | generate-test-checks.py --source foo.mlir -i --source_delim_regex='gpu.func @'15 16The script will heuristically generate CHECK/CHECK-LABEL commands for each line17within the file. By default this script will also try to insert string18substitution blocks for all SSA value names. If --source file is specified, the19script will attempt to insert the generated CHECKs to the source file by looking20for line positions matched by --source_delim_regex.21 22The script is designed to make adding checks to a test case fast, it is *not*23designed to be authoritative about what constitutes a good test!24"""25 26# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.27# See https://llvm.org/LICENSE.txt for license information.28# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception29 30import argparse31import os # Used to advertise this file's name ("autogenerated_note").32import re33import sys34from collections import Counter35 36ADVERT_BEGIN = "// NOTE: Assertions have been autogenerated by "37ADVERT_END = """38// This script is intended to make adding checks to a test case quick and easy.39// It is *not* authoritative about what constitutes a good test. After using the40// script, be sure to review and refine the generated checks. For example,41// CHECK lines should be minimized and named to reflect the test’s intent.42// For comprehensive guidelines, see:43// * https://mlir.llvm.org/getting_started/TestingGuide/44"""45 46 47# Regex command to match an SSA identifier.48SSA_RE_STR = "[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*"49SSA_RE = re.compile(SSA_RE_STR)50 51# Regex matching `dialect.op_name` (e.g. `vector.transfer_read`).52SSA_OP_NAME_RE = re.compile(r"\b(?:\s=\s[a-z_]+)[.]([a-z_]+)\b")53 54# Regex matching the left-hand side of an assignment55SSA_RESULTS_STR = r'\s*(%' + SSA_RE_STR + r')(\s*,\s*(%' + SSA_RE_STR + r'))*\s*='56SSA_RESULTS_RE = re.compile(SSA_RESULTS_STR)57 58# Regex matching attributes59ATTR_RE_STR = r'(#[a-zA-Z._-][a-zA-Z0-9._-]*)'60ATTR_RE = re.compile(ATTR_RE_STR)61 62# Regex matching the left-hand side of an attribute definition63ATTR_DEF_RE_STR = r'\s*' + ATTR_RE_STR + r'\s*='64ATTR_DEF_RE = re.compile(ATTR_DEF_RE_STR)65 66 67# Class used to generate and manage string substitution blocks for SSA value68# names.69class VariableNamer:70 def __init__(self, variable_names):71 self.scopes = []72 # Counter for generic FileCHeck names, e.g. VAL_#N73 self.name_counter = 074 # Counters for FileCheck names derived from Op names, e.g.75 # TRANSFER_READ_#N (based on `vector.transfer_read`). Note, there's a76 # dedicated counter for every Op type present in the input.77 self.op_name_counter = Counter()78 79 # Number of variable names to still generate in parent scope80 self.generate_in_parent_scope_left = 081 82 # Parse variable names83 self.variable_names = [name.upper() for name in variable_names.split(',')]84 self.used_variable_names = set()85 86 # Generate the following 'n' variable names in the parent scope.87 def generate_in_parent_scope(self, n):88 self.generate_in_parent_scope_left = n89 90 # Generate a substitution name for the given ssa value name.91 def generate_name(self, source_variable_name, use_ssa_name, op_name=""):92 93 # Compute variable name94 variable_name = (95 self.variable_names.pop(0) if len(self.variable_names) > 0 else ""96 )97 if variable_name == "":98 # If `use_ssa_name` is set, use the MLIR SSA value name to generate99 # a FileCHeck substation string. As FileCheck requires these100 # strings to start with a character, skip MLIR variables starting101 # with a digit (e.g. `%0`).102 #103 # The next fallback option is to use the op name, if the104 # corresponding match succeeds.105 #106 # If neither worked, use a generic name: `VAL_#N`.107 if use_ssa_name and source_variable_name[0].isalpha():108 variable_name = source_variable_name.upper()109 elif op_name != "":110 variable_name = (111 op_name.upper() + "_" + str(self.op_name_counter[op_name])112 )113 self.op_name_counter[op_name] += 1114 else:115 variable_name = "VAL_" + str(self.name_counter)116 self.name_counter += 1117 118 # Scope where variable name is saved119 scope = len(self.scopes) - 1120 if self.generate_in_parent_scope_left > 0:121 self.generate_in_parent_scope_left -= 1122 scope = len(self.scopes) - 2123 assert(scope >= 0)124 125 # Save variable126 if variable_name in self.used_variable_names:127 raise RuntimeError(variable_name + ': duplicate variable name')128 self.scopes[scope][source_variable_name] = variable_name129 self.used_variable_names.add(variable_name)130 131 return variable_name132 133 # Push a new variable name scope.134 def push_name_scope(self):135 self.scopes.append({})136 137 # Pop the last variable name scope.138 def pop_name_scope(self):139 self.scopes.pop()140 141 # Return the level of nesting (number of pushed scopes).142 def num_scopes(self):143 return len(self.scopes)144 145 # Reset the counter and used variable names.146 def clear_names(self):147 self.name_counter = 0148 self.used_variable_names = set()149 self.op_name_counter.clear()150 151class AttributeNamer:152 153 def __init__(self, attribute_names):154 self.name_counter = 0155 self.attribute_names = [name.upper() for name in attribute_names.split(',')]156 self.map = {}157 self.used_attribute_names = set()158 159 # Generate a substitution name for the given attribute name.160 def generate_name(self, source_attribute_name):161 162 # Compute FileCheck name163 attribute_name = self.attribute_names.pop(0) if len(self.attribute_names) > 0 else ''164 if attribute_name == '':165 attribute_name = "ATTR_" + str(self.name_counter)166 self.name_counter += 1167 168 # Prepend global symbol169 attribute_name = '$' + attribute_name170 171 # Save attribute172 if attribute_name in self.used_attribute_names:173 raise RuntimeError(attribute_name + ': duplicate attribute name')174 self.map[source_attribute_name] = attribute_name175 self.used_attribute_names.add(attribute_name)176 return attribute_name177 178 # Get the saved substitution name for the given attribute name. If no name179 # has been generated for the given attribute yet, None is returned.180 def get_name(self, source_attribute_name):181 return self.map.get(source_attribute_name)182 183# Return the number of SSA results in a line of type184# %0, %1, ... = ...185# The function returns 0 if there are no results.186def get_num_ssa_results(input_line):187 m = SSA_RESULTS_RE.match(input_line)188 return m.group().count('%') if m else 0189 190 191# Process a line of input that has been split at each SSA identifier '%'.192def process_line(line_chunks, variable_namer, use_ssa_name=False, strict_name_re=False):193 output_line = ""194 195 # Process the rest that contained an SSA value name.196 for chunk in line_chunks:197 ssa = SSA_RE.match(chunk)198 op_name_with_dialect = SSA_OP_NAME_RE.search(chunk)199 ssa_name = ssa.group(0) if ssa is not None else ""200 op_name = (201 op_name_with_dialect.group(1) if op_name_with_dialect is not None else ""202 )203 204 # Check if an existing variable exists for this name.205 variable = None206 for scope in variable_namer.scopes:207 variable = scope.get(ssa_name)208 if variable is not None:209 break210 211 # If one exists, then output the existing name.212 if variable is not None:213 output_line += "%[[" + variable + "]]"214 else:215 # Otherwise, generate a new variable.216 variable = variable_namer.generate_name(ssa_name, use_ssa_name, op_name)217 if strict_name_re:218 # Use stricter regexp for the variable name, if requested.219 # Greedy matching may cause issues with the generic '.*'220 # regexp when the checks are split across several221 # lines (e.g. for CHECK-SAME).222 output_line += "%[[" + variable + ":" + SSA_RE_STR + "]]"223 else:224 output_line += "%[[" + variable + ":.*]]"225 226 # Append the non named group.227 output_line += chunk[len(ssa_name) :]228 229 return output_line.rstrip() + "\n"230 231 232# Process the source file lines. The source file doesn't have to be .mlir.233def process_source_lines(source_lines, args):234 source_split_re = re.compile(args.source_delim_regex)235 236 source_segments = [[]]237 for line in source_lines:238 # Remove previous CHECK lines.239 if line.find(args.check_prefix) != -1:240 continue241 # Segment the file based on --source_delim_regex.242 if source_split_re.search(line):243 source_segments.append([])244 245 source_segments[-1].append(line + "\n")246 return source_segments247 248 249def process_attribute_definition(line, attribute_namer):250 m = ATTR_DEF_RE.match(line)251 if m:252 attribute_name = attribute_namer.generate_name(m.group(1))253 return (254 "// CHECK: #[["255 + attribute_name256 + ":.+]] ="257 # The rest of the line may contain attribute references,258 # so we have to process them.259 + process_attribute_references(line[len(m.group(0)) :], attribute_namer)260 + "\n"261 )262 return None263 264def process_attribute_references(line, attribute_namer):265 266 output_line = ''267 components = ATTR_RE.split(line)268 for component in components:269 m = ATTR_RE.match(component)270 attribute_name = attribute_namer.get_name(m.group(1)) if m else None271 if attribute_name:272 output_line += f"#[[{attribute_name}]]{component[len(m.group()):]}"273 else:274 output_line += component275 return output_line276 277# Pre-process a line of input to remove any character sequences that will be278# problematic with FileCheck.279def preprocess_line(line):280 # Replace any `{{` with escaped replacements. `{{` corresponds to regex281 # checks in FileCheck.282 output_line = line.replace("{{", "{{\\{\\{}}")283 284 # Replace any double brackets, '[[' with escaped replacements. '[['285 # corresponds to variable names in FileCheck.286 output_line = output_line.replace("[[", "{{\\[\\[}}")287 288 # Replace any single brackets that are followed by an SSA identifier, the289 # identifier will be replace by a variable; Creating the same situation as290 # above.291 output_line = output_line.replace("[%", "{{\\[}}%")292 293 return output_line294 295 296def main():297 parser = argparse.ArgumentParser(298 description=__doc__, formatter_class=argparse.RawTextHelpFormatter299 )300 parser.add_argument(301 "--check-prefix", default="CHECK", help="Prefix to use from check file."302 )303 parser.add_argument(304 "-o", "--output", nargs="?", type=argparse.FileType("w"), default=None305 )306 parser.add_argument(307 "input", nargs="?", type=argparse.FileType("r"), default=sys.stdin308 )309 parser.add_argument(310 "--source",311 type=str,312 help="Print each CHECK chunk before each delimeter line in the source"313 "file, respectively. The delimeter lines are identified by "314 "--source_delim_regex.",315 )316 parser.add_argument("--source_delim_regex", type=str, default="func @")317 parser.add_argument(318 "--starts_from_scope",319 type=int,320 default=1,321 help="Omit the top specified level of content. For example, by default "322 'it omits "module {"',323 )324 parser.add_argument("-i", "--inplace", action="store_true", default=False)325 parser.add_argument(326 "--variable_names",327 type=str,328 default='',329 help="Names to be used in FileCheck regular expression to represent SSA "330 "variables in the order they are encountered. Separate names with commas, "331 "and leave empty entries for default names (e.g.: 'DIM,,SUM,RESULT')")332 parser.add_argument(333 "--attribute_names",334 type=str,335 default='',336 help="Names to be used in FileCheck regular expression to represent "337 "attributes in the order they are defined. Separate names with commas,"338 "commas, and leave empty entries for default names (e.g.: 'MAP0,,,MAP1')")339 parser.add_argument(340 "--strict_name_re",341 type=bool,342 default=False,343 help="Set to true to use stricter regex for CHECK-SAME directives. "344 "Use when Greedy matching causes issues with the generic '.*'",345 )346 347 args = parser.parse_args()348 349 # Open the given input file.350 input_lines = [l.rstrip() for l in args.input]351 args.input.close()352 353 # Generate a note used for the generated check file.354 script_name = os.path.basename(__file__)355 autogenerated_note = ADVERT_BEGIN + "utils/" + script_name + "\n" + ADVERT_END356 357 source_segments = None358 if args.source:359 with open(args.source, "r") as f:360 raw_source = f.read().replace(autogenerated_note, "")361 raw_source_lines = [l.rstrip() for l in raw_source.splitlines()]362 source_segments = process_source_lines(raw_source_lines, args)363 364 if args.inplace:365 assert args.output is None366 output = open(args.source, "w")367 elif args.output is None:368 output = sys.stdout369 else:370 output = args.output371 372 output_segments = [[]]373 374 # Namers375 variable_namer = VariableNamer(args.variable_names)376 attribute_namer = AttributeNamer(args.attribute_names)377 378 # Store attribute definitions to emit at appropriate scope379 pending_attr_defs = []380 381 # Process lines382 for input_line in input_lines:383 if not input_line:384 continue385 386 # When using `--starts_from_scope=0` to capture module lines, the file387 # split needs to be skipped, otherwise a `CHECK: // -----` is inserted.388 if input_line.startswith("// -----"):389 continue390 391 if ATTR_DEF_RE.match(input_line):392 pending_attr_defs.append(input_line)393 continue394 395 # Lines with blocks begin with a ^. These lines have a trailing comment396 # that needs to be stripped.397 lstripped_input_line = input_line.lstrip()398 is_block = lstripped_input_line[0] == "^"399 if is_block:400 input_line = input_line.rsplit("//", 1)[0].rstrip()401 402 cur_level = variable_namer.num_scopes()403 404 # If the line starts with a '}', pop the last name scope.405 if lstripped_input_line[0] == "}":406 variable_namer.pop_name_scope()407 cur_level = variable_namer.num_scopes()408 409 # If the line ends with a '{', push a new name scope.410 if input_line[-1] == "{":411 variable_namer.push_name_scope()412 if cur_level == args.starts_from_scope:413 output_segments.append([])414 415 # Result SSA values must still be pushed to parent scope416 num_ssa_results = get_num_ssa_results(input_line)417 variable_namer.generate_in_parent_scope(num_ssa_results)418 419 # Omit lines at the near top level e.g. "module {".420 if cur_level < args.starts_from_scope:421 continue422 423 if len(output_segments[-1]) == 0:424 variable_namer.clear_names()425 426 # Preprocess the input to remove any sequences that may be problematic with427 # FileCheck.428 input_line = preprocess_line(input_line)429 430 # Process uses of attributes in this line431 input_line = process_attribute_references(input_line, attribute_namer)432 433 # Split the line at the each SSA value name.434 ssa_split = input_line.split("%")435 436 # If this is a top-level operation use 'CHECK-LABEL', otherwise 'CHECK:'.437 if len(output_segments[-1]) != 0 or not ssa_split[0]:438 output_line = "// " + args.check_prefix + ": "439 # Pad to align with the 'LABEL' statements.440 output_line += " " * len("-LABEL")441 442 # Output the first line chunk that does not contain an SSA name.443 output_line += ssa_split[0]444 445 # Process the rest of the input line.446 output_line += process_line(ssa_split[1:], variable_namer)447 448 else:449 # Emit any pending attribute definitions at the start of this scope450 for attr in pending_attr_defs:451 attr_line = process_attribute_definition(attr, attribute_namer)452 if attr_line:453 output_segments[-1].append(attr_line)454 pending_attr_defs.clear()455 456 # Output the first line chunk that does not contain an SSA name for the457 # label.458 output_line = "// " + args.check_prefix + "-LABEL: " + ssa_split[0] + "\n"459 460 # Process the rest of the input line on separate check lines.461 for argument in ssa_split[1:]:462 output_line += "// " + args.check_prefix + "-SAME: "463 464 # Pad to align with the original position in the line (i.e. where the label ends),465 # unless the label is more than 20 chars long, in which case pad with 4 spaces466 # (this is to avoid deep indentation).467 label_length = len(ssa_split[0])468 pad_depth = label_length if label_length < 21 else 4469 output_line += " " * pad_depth470 471 # Process the rest of the line. Use the original SSA name to generate the LIT472 # variable names.473 use_ssa_names = True474 output_line += process_line(475 [argument], variable_namer, use_ssa_names, args.strict_name_re476 )477 478 # Append the output line.479 output_segments[-1].append(output_line)480 481 output.write(autogenerated_note + "\n")482 483 # Write the output.484 if source_segments:485 assert len(output_segments) == len(source_segments)486 for check_segment, source_segment in zip(output_segments, source_segments):487 for line in check_segment:488 output.write(line)489 for line in source_segment:490 output.write(line)491 else:492 for segment in output_segments:493 output.write("\n")494 for output_line in segment:495 output.write(output_line)496 output.write("\n")497 output.close()498 499 500if __name__ == "__main__":501 main()502