v3.9 update (#2203)
* v3.9 update * voidD --------- Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
GitHub
356
tools/util/scripts/split_test_cmake.py
Normal file
356
tools/util/scripts/split_test_cmake.py
Normal file
@ -0,0 +1,356 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
|
||||
"""
|
||||
Given a set of test files to be included in a CMake target, this script extracts
|
||||
the TEST definitions from each file, writes them into new files, and prints the names
|
||||
of the new files so that they can be processed as part of a new CMake target.
|
||||
|
||||
For example, given a set of --src_files test_a.cu test_b.cu containing 3 and 2 TEST
|
||||
definitions, respectively, this script would produce:
|
||||
test_a_000.cu
|
||||
test_a_001.cu
|
||||
test_a_002.cu
|
||||
test_b_000.cu
|
||||
test_b_001.cu
|
||||
|
||||
The splitting follows a fairly rudimentary algorithm that does not support all valid C++ programs.
|
||||
We walk through a given input test file line by line. Any lines that are not within a TEST definition is added to a running
|
||||
"filler" text. When a TEST definition is encountered, the current filler text becomes the prefix
|
||||
for that test. All subsequent lines are considered to be part of the TEST definition until the
|
||||
number of starting function braces ('{') match the number of closing function braces ('}'). When
|
||||
these counts are equal, the TEST definition is considered to be completed. At this point, we return
|
||||
to adding lines to the "filler" text until a new TEST definition is encountered. Any "filler" text
|
||||
following a TEST definition is added to the suffix of that TEST definition (this is useful for finishing
|
||||
off #if statements, as is common in unit tests.).
|
||||
|
||||
A state machine illustrating this algorithm at a high level is provided in the source below.
|
||||
|
||||
Example: Suppose an input test `test.cu` has the following source:
|
||||
// COPYRIGHT
|
||||
#include <iostream>
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
// Test #1
|
||||
TEST(SM90_a, 256x128x64_2x2x1) {
|
||||
std::cout << "Test #1" << std::endl;
|
||||
}
|
||||
|
||||
// Test #2
|
||||
TEST(SM90_b, 256x128x64_1x1x1) {
|
||||
std::cout << "Test #2" << std::endl;
|
||||
}
|
||||
|
||||
#endif defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
The contents of the two resulting test files will be:
|
||||
$ cat test_000.cu
|
||||
// COPYRIGHT
|
||||
#include <iostream>
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
// Test #1
|
||||
TEST(SM90_a, 256x128x64_2x2x1) {
|
||||
std::cout << "Test #1" << std::endl;
|
||||
}
|
||||
|
||||
// Test #2
|
||||
|
||||
#endif defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
$ cat test_001.cu
|
||||
// COPYRIGHT
|
||||
#include <iostream>
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
// Test #1
|
||||
|
||||
// Test #2
|
||||
TEST(SM90_b, 256x128x64_1x1x1) {
|
||||
std::cout << "Test #2" << std::endl;
|
||||
}
|
||||
|
||||
#endif defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
Notice that each of test_000.cu and test_001.cu contain comments that appear outside
|
||||
the TEST definitions not included in each file. This is by design, as these
|
||||
would be considered "filler" text.
|
||||
|
||||
As expected, some cases can't be handled. Below is a non-exhaustive list:
|
||||
1. New TEST following the closing '}' of a TEST case on the same line:
|
||||
TEST(x, y) {
|
||||
// Do stuff
|
||||
} TEST(a, b) {
|
||||
|
||||
In this case, "TEST(a, b) {" will be ignored
|
||||
|
||||
2. Preprocessor macros that occur midway through a test case and extend
|
||||
beyond the conclusion of a testcase
|
||||
|
||||
Example:
|
||||
TEST(a, b) {
|
||||
// Do stuff
|
||||
#if X
|
||||
// Do more stuff
|
||||
}
|
||||
#else
|
||||
// Do other stuff
|
||||
}
|
||||
#endif
|
||||
"""
|
||||
|
||||
|
||||
import argparse
|
||||
import enum
|
||||
import os
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("cmake_target", type=str,
|
||||
help="Name of the CMake target being generated.")
|
||||
parser.add_argument("src_dir", type=str,
|
||||
help="Path to the directory containing test files.")
|
||||
parser.add_argument("--src_files", nargs='+',
|
||||
help="Files containing TEST instances to split.")
|
||||
parser.add_argument("--max_tests_per_file", type=int, default=1,
|
||||
help="Maximum number of TEST instances per file.")
|
||||
parser.add_argument("--dst_dir", type=str,
|
||||
help="Path to the directory to which to write new test files. If not set, uses src_dir.")
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
if args.dst_dir == None:
|
||||
args.dst_dir = args.src_dir
|
||||
|
||||
|
||||
class Testcase:
|
||||
"""
|
||||
Lightweight tracker of test-case processing status
|
||||
"""
|
||||
def __init__(self, prefix_text):
|
||||
# Any text that preceded the TEST definition that was
|
||||
# not part of another TEST definition
|
||||
self.prefix = prefix_text
|
||||
|
||||
# Any text within the TEST definition
|
||||
self.test = ""
|
||||
|
||||
# Any text that follows the completion of the TEST definition
|
||||
# and is not included in other TEST definitions
|
||||
self.suffix = ""
|
||||
|
||||
# Whether the test's definition has concluded
|
||||
self.completed = False
|
||||
|
||||
# Current balance of opening and closing curly brackets in
|
||||
# the TEST definition. '{' increments the count and '}' decrements it.
|
||||
# A value of 0 (when self.completed == False) indicates that the test
|
||||
# has completed.
|
||||
self.curly_bracket_balance = 0
|
||||
|
||||
|
||||
class ParseState(enum.Enum):
|
||||
"""
|
||||
State machine for processing.
|
||||
Transitions occur on each line encountered in the soruce file
|
||||
|
||||
|
||||
Line does not contain 'TEST('
|
||||
+----+
|
||||
| |
|
||||
| v 'TEST('
|
||||
+--------+ encountered +--------------------------+
|
||||
------>| Filler | -----------------------> | TestDeclaredWaitingStart |
|
||||
+--------+ +--------------------------+
|
||||
^ |
|
||||
Number of '{' | | First '{' encountered
|
||||
equals number of | +--------+ |
|
||||
'}' encountered +-----------| InTest | <------------------+
|
||||
+--------+
|
||||
| ^
|
||||
| |
|
||||
+----+
|
||||
Number of '{' encountered
|
||||
exceeds number of '}' encountered
|
||||
"""
|
||||
|
||||
|
||||
# Any text that is not part of a TEST case
|
||||
Filler = 0
|
||||
|
||||
# Processing text within the first { of the TEST case
|
||||
# and before the en of the final } of the TEST case
|
||||
InTest = 1
|
||||
|
||||
# Processing text from the start of the TEST definition
|
||||
# but before the first {. This could occur if the opening {
|
||||
# occurs on a separate line than the TEST definition.
|
||||
TestDeclaredWaitingStart = 2
|
||||
|
||||
|
||||
cmake_src_list = []
|
||||
for filename in args.src_files:
|
||||
if '.' not in filename:
|
||||
# Add any non-filename arguments to the command list by default
|
||||
cmake_src_list.append(filename)
|
||||
continue
|
||||
|
||||
if '/' in filename:
|
||||
raise Exception(
|
||||
f"Source files passed to {__file__} must be within the same directory "
|
||||
"as the CMakeLists defining the target using the files. "
|
||||
f"Provided path {filename} is in a different directory.")
|
||||
|
||||
full_filename = os.path.join(args.src_dir, filename)
|
||||
with open(full_filename, 'r') as infile:
|
||||
lines = infile.readlines()
|
||||
|
||||
# Find the number of instances of "TEST("
|
||||
ntest = sum([1 for line in lines if "TEST(" in line])
|
||||
|
||||
if ntest <= args.max_tests_per_file:
|
||||
# File contains fewer than max_tests_per_file TEST instances. It does
|
||||
# not need to be split
|
||||
cmake_src_list.append(filename)
|
||||
continue
|
||||
|
||||
# Current state of the parsing state machine. We start with filler text
|
||||
state = ParseState.Filler
|
||||
|
||||
# List of individual TESTs found
|
||||
tests = []
|
||||
|
||||
# Ongoing text that is not included in a TEST definition. This will serve
|
||||
# as the prefix for any yet-to-be encountered TEST definitions.
|
||||
filler_text = ""
|
||||
|
||||
def add_filler_text(text):
|
||||
global filler_text
|
||||
# Add new text to the ongoing filler text and to the suffixes of
|
||||
# any completed tests
|
||||
filler_text += text
|
||||
for i in range(len(tests)):
|
||||
if tests[i].completed:
|
||||
tests[i].suffix += text
|
||||
|
||||
for line in lines:
|
||||
if state == ParseState.Filler:
|
||||
# We are not currently within a TEST definition.
|
||||
|
||||
if 'TEST(' in line:
|
||||
# We have encountered a new TEST( case. Any text preceding this
|
||||
# must be added to the filler text (e.g., if we have a line of the form:
|
||||
# "static constexpr int Val = 4; TEST(blah) {"
|
||||
# then "static constexpr int Val = 4;" needs to be included in filler
|
||||
# text, as it could be used by subsequent tests.)
|
||||
splits = line.split('TEST')
|
||||
|
||||
# There should not be more than one TEST definition on a given line
|
||||
assert len(splits) <= 2
|
||||
|
||||
if len(splits) > 1:
|
||||
if not splits[0].isspace():
|
||||
# Only add text to filler if there are non-whitespace charcters
|
||||
# preceding the TEST definition in the line
|
||||
filler_text += splits[0]
|
||||
|
||||
# The new line is just the TEST-related line
|
||||
line = 'TEST' + splits[-1]
|
||||
|
||||
# Add tests and transtion to TestDeclaredWaitingStart state.
|
||||
# Do not add the line to the test text of the new test case; this
|
||||
# will be done in either the TestDeclaredWaitingStart state processing
|
||||
# below or in the InTest state processing below.
|
||||
tests.append(Testcase(filler_text))
|
||||
state = ParseState.TestDeclaredWaitingStart
|
||||
else:
|
||||
# Any remaining filler text is added to the running filler_text
|
||||
# which will be used as the prefix for any new tests, and to the
|
||||
# suffix of any completed tests
|
||||
add_filler_text(line)
|
||||
|
||||
if state == ParseState.TestDeclaredWaitingStart:
|
||||
# We have seen a TEST definition but have not yet seen its opening {.
|
||||
|
||||
if '{' in line:
|
||||
# The first curly bracket for the TEST definition has been found.
|
||||
# Advance to state InTests. Do not add the line to the test's text
|
||||
# or change the curly-brace balance of the test; these will be done
|
||||
# when processing the state == ParseState.InTest condition below.
|
||||
state = ParseState.InTest
|
||||
else:
|
||||
tests[-1].test += line
|
||||
|
||||
if state == ParseState.InTest:
|
||||
# We are currently within a TEST definition.
|
||||
# Process lines character-by-character looking for opening and closing
|
||||
# braces. If we reach parity between opening and closing braces, the
|
||||
# test is considered done.
|
||||
filler_text_to_add = ""
|
||||
for char in line:
|
||||
if not tests[-1].completed:
|
||||
tests[-1].test += char
|
||||
if char == '{':
|
||||
tests[-1].curly_bracket_balance += 1
|
||||
elif char == '}':
|
||||
tests[-1].curly_bracket_balance -= 1
|
||||
if tests[-1].curly_bracket_balance == 0:
|
||||
tests[-1].completed = True
|
||||
else:
|
||||
filler_text_to_add += char
|
||||
|
||||
if filler_text_to_add != "" and (not filler_text_to_add.isspace() or '\n' in filler_text_to_add):
|
||||
add_filler_text('\n' + filler_text_to_add)
|
||||
|
||||
if tests[-1].completed:
|
||||
state = ParseState.Filler
|
||||
|
||||
# Write out the new files for tests
|
||||
filename_prefix, filename_suffix = filename.split('.')
|
||||
for i, test in enumerate(tests):
|
||||
assert test.completed
|
||||
new_filename = filename_prefix + '_' + str(i).zfill(3) + '.' + filename_suffix
|
||||
full_new_filename = os.path.join(args.dst_dir, new_filename)
|
||||
|
||||
# Replace any '\' with '/'. CMake doesn't like '\'.
|
||||
full_new_filename = full_new_filename.replace('\\', '/')
|
||||
|
||||
with open(full_new_filename, 'w') as outfile:
|
||||
outfile.write(test.prefix + test.test + test.suffix)
|
||||
cmake_src_list.append(full_new_filename)
|
||||
|
||||
|
||||
for cmake_file in cmake_src_list:
|
||||
print(cmake_file)
|
||||
Reference in New Issue
Block a user