3131import os # Used to advertise this file's name ("autogenerated_note").
3232import re
3333import sys
34+ from collections import Counter
3435
3536ADVERT_BEGIN = "// NOTE: Assertions have been autogenerated by "
3637ADVERT_END = """
37- // The script is designed to make adding checks to
38- // a test case fast, it is *not* designed to be authoritative
39- // about what constitutes a good test! The CHECK should be
40- // minimized and named to reflect the test intent.
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 the
40+ // 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/
4144"""
4245
4346
4447# Regex command to match an SSA identifier.
4548SSA_RE_STR = "[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*"
4649SSA_RE = re .compile (SSA_RE_STR )
4750
51+ # Regex matching `dialect.op_name` (e.g. `vector.transfer_read`).
52+ SSA_OP_NAME_RE = re .compile (r"\b(?:\s=\s[a-z_]+)[.]([a-z_]+)\b" )
53+
4854# Regex matching the left-hand side of an assignment
4955SSA_RESULTS_STR = r'\s*(%' + SSA_RE_STR + r')(\s*,\s*(%' + SSA_RE_STR + r'))*\s*='
5056SSA_RESULTS_RE = re .compile (SSA_RESULTS_STR )
6369class VariableNamer :
6470 def __init__ (self , variable_names ):
6571 self .scopes = []
72+ # Counter for generic FileCHeck names, e.g. VAL_#N
6673 self .name_counter = 0
74+ # Counters for FileCheck names derived from Op names, e.g.
75+ # TRANSFER_READ_#N (based on `vector.transfer_read`). Note, there's a
76+ # dedicated counter for every Op type present in the input.
77+ self .op_name_counter = Counter ()
6778
6879 # Number of variable names to still generate in parent scope
6980 self .generate_in_parent_scope_left = 0
@@ -77,17 +88,29 @@ def generate_in_parent_scope(self, n):
7788 self .generate_in_parent_scope_left = n
7889
7990 # Generate a substitution name for the given ssa value name.
80- def generate_name (self , source_variable_name , use_ssa_name ):
91+ def generate_name (self , source_variable_name , use_ssa_name , op_name = "" ):
8192
8293 # Compute variable name
83- variable_name = self .variable_names .pop (0 ) if len (self .variable_names ) > 0 else ''
84- if variable_name == '' :
94+ variable_name = (
95+ self .variable_names .pop (0 ) if len (self .variable_names ) > 0 else ""
96+ )
97+ if variable_name == "" :
8598 # If `use_ssa_name` is set, use the MLIR SSA value name to generate
8699 # a FileCHeck substation string. As FileCheck requires these
87100 # strings to start with a character, skip MLIR variables starting
88101 # with a digit (e.g. `%0`).
102+ #
103+ # The next fallback option is to use the op name, if the
104+ # corresponding match succeeds.
105+ #
106+ # If neither worked, use a generic name: `VAL_#N`.
89107 if use_ssa_name and source_variable_name [0 ].isalpha ():
90108 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 ] += 1
91114 else :
92115 variable_name = "VAL_" + str (self .name_counter )
93116 self .name_counter += 1
@@ -123,6 +146,7 @@ def num_scopes(self):
123146 def clear_names (self ):
124147 self .name_counter = 0
125148 self .used_variable_names = set ()
149+ self .op_name_counter .clear ()
126150
127151class AttributeNamer :
128152
@@ -170,8 +194,12 @@ def process_line(line_chunks, variable_namer, use_ssa_name=False, strict_name_re
170194
171195 # Process the rest that contained an SSA value name.
172196 for chunk in line_chunks :
173- m = SSA_RE .match (chunk )
174- ssa_name = m .group (0 ) if m is not None else ''
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+ )
175203
176204 # Check if an existing variable exists for this name.
177205 variable = None
@@ -185,7 +213,7 @@ def process_line(line_chunks, variable_namer, use_ssa_name=False, strict_name_re
185213 output_line += "%[[" + variable + "]]"
186214 else :
187215 # Otherwise, generate a new variable.
188- variable = variable_namer .generate_name (ssa_name , use_ssa_name )
216+ variable = variable_namer .generate_name (ssa_name , use_ssa_name , op_name )
189217 if strict_name_re :
190218 # Use stricter regexp for the variable name, if requested.
191219 # Greedy matching may cause issues with the generic '.*'
0 commit comments