Skip to content

Commit 7361395

Browse files
author
Kshitij
committed
[mlir][tblgen] Adds support for embedded LIT tests in TableGen records
Introduces a new Testable base class that allows TableGen records (starting with Pass records) to embed LIT test definitions directly within their definitions. This enables co-locating tests with pass definitions for better maintainability. Key components: - Testable.td: Base class for records that can have embedded tests - LitTestGen.cpp: TableGen backend to extract and generate LIT test files - AddMLIR.cmake: CMake function to process embedded tests with usage examples - PassBase.td: Updated Pass class to extend Testable Usage example in CMake: add_embedded_lit_tests(MyPassesEmbeddedTests ${CMAKE_CURRENT_SOURCE_DIR}/include/MyPasses.td ${CMAKE_CURRENT_SOURCE_DIR}/test/Passes/) # Add LIT test generation target as a dependency to some # other target add_library(someLib DEPENDS MyPassesEmbeddedTests)
1 parent 98f1ae0 commit 7361395

File tree

6 files changed

+379
-1
lines changed

6 files changed

+379
-1
lines changed

mlir/cmake/modules/AddMLIR.cmake

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,3 +762,103 @@ function(mlir_target_link_libraries target type)
762762
target_link_libraries(${target} ${type} ${ARGN})
763763
endif()
764764
endfunction()
765+
766+
# Extracts LIT tests embedded in `Testable` records in `tblgen_file`
767+
# and generates a file per test in `output_dir`
768+
#
769+
# Example usage:
770+
# # Extract tests from MyPasses.td and generate them in test/Passes/
771+
# add_embedded_lit_tests(MyPassesEmbeddedTests
772+
# ${CMAKE_CURRENT_SOURCE_DIR}/include/MyPasses.td
773+
# ${CMAKE_CURRENT_SOURCE_DIR}/test/Passes/)
774+
#
775+
# # This will:
776+
# # 1. Process MyPasses.td with mlir-tblgen --gen-lit-tests
777+
# # 2. Extract individual test files to test/Passes/
778+
# # 3. Generate files like: test/Passes/generated_MyPass_test1.mlir
779+
#
780+
function(add_embedded_lit_tests target tblgen_file output_dir)
781+
set(LLVM_TARGET_DEFINITIONS ${tblgen_file})
782+
783+
# Extraction script content
784+
set(EXTRACT_SCRIPT_CONTENT [[
785+
# Generated extraction script
786+
if(NOT CONSOLIDATED_FILE)
787+
message(FATAL_ERROR "CONSOLIDATED_FILE variable is required")
788+
endif()
789+
790+
if(NOT OUTPUT_DIR)
791+
message(FATAL_ERROR "OUTPUT_DIR variable is required")
792+
endif()
793+
794+
if(NOT EXISTS ${CONSOLIDATED_FILE})
795+
message(FATAL_ERROR "Consolidated file does not exist: ${CONSOLIDATED_FILE}")
796+
endif()
797+
798+
# Read the consolidated file
799+
file(READ ${CONSOLIDATED_FILE} file_content)
800+
801+
# Split into lines for processing
802+
string(REPLACE "\n" ";" lines "${file_content}")
803+
804+
set(current_filename "")
805+
set(current_content "")
806+
set(in_test_block FALSE)
807+
set(extracted_test_files)
808+
809+
foreach(line IN LISTS lines)
810+
# Check for filename line
811+
if(line MATCHES "^// File: (.+)$")
812+
set(current_filename "${CMAKE_MATCH_1}")
813+
endif()
814+
815+
# Check for BEGIN marker
816+
if(line MATCHES "^// --- BEGIN .+ ---$")
817+
set(in_test_block TRUE)
818+
set(current_content "")
819+
# Check for END marker
820+
elseif(line MATCHES "^// --- END .+ ---$")
821+
set(in_test_block FALSE)
822+
823+
# Write the extracted content to file
824+
if(current_filename AND current_content)
825+
file(MAKE_DIRECTORY ${OUTPUT_DIR})
826+
file(WRITE ${OUTPUT_DIR}/${current_filename} "${current_content}")
827+
message(STATUS "Extracted test file: ${current_filename}")
828+
list(APPEND extracted_test_files ${current_filename})
829+
endif()
830+
831+
set(current_filename "")
832+
set(current_content "")
833+
# Collect content within BEGIN/END block
834+
elseif(in_test_block)
835+
string(APPEND current_content "${line}\n")
836+
endif()
837+
endforeach()
838+
839+
list(LENGTH extracted_test_files num_extracted_files)
840+
message(STATUS "Extracted ${num_extracted_files} test files to ${OUTPUT_DIR}")
841+
]])
842+
843+
# Write extraction script to a file in the build directory
844+
file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/extract_lit_tests.cmake "${EXTRACT_SCRIPT_CONTENT}")
845+
846+
# Process tblgen_file and generate a file with all embedded LIT
847+
# tests in tblgen_file
848+
get_filename_component(tblgen_name ${tblgen_file} NAME_WE)
849+
set(consolidated_output_file ${tblgen_name}_extracted_lit_tests.txt)
850+
mlir_tablegen(${consolidated_output_file} --gen-lit-tests)
851+
852+
# Add public tablegen target to trigger builds on changes in tblgen_file
853+
add_public_tablegen_target(${target})
854+
855+
# Call the extraction script to extract all LIT tests into individual
856+
# `.mlir` test files
857+
add_custom_command(TARGET ${target} POST_BUILD
858+
COMMAND ${CMAKE_COMMAND}
859+
-DCONSOLIDATED_FILE=${CMAKE_CURRENT_BINARY_DIR}/${consolidated_output_file}
860+
-DOUTPUT_DIR=${output_dir}
861+
-P ${CMAKE_CURRENT_BINARY_DIR}/extract_lit_tests.cmake
862+
COMMENT "Extracting LIT tests to individual files"
863+
)
864+
endfunction()

mlir/include/mlir/IR/Testable.td

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
//===-- Testable.td - Testable type definition file --------*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file contains the definition of the `Testable` type.
10+
//
11+
// Any type whose records can have corresponding LIT tests (eg - Pass) can extend
12+
// `Testable` in order to be able to embed LIT tests within record definitions.
13+
//
14+
//===----------------------------------------------------------------------===//
15+
16+
#ifndef TESTABLE
17+
#define TESTABLE
18+
19+
// Represents a LIT test record in TableGen
20+
class LitTest<string name, code snippet, list<string> run = [], list<string> check = []> {
21+
// The name of the generated test file
22+
string testFileName = name;
23+
24+
// The IR snippet/code to be tested
25+
code irSnippet = snippet;
26+
27+
// The RUN commands for the test (e.g., "mlir-opt %s")
28+
list<string> runLines = run;
29+
30+
// Expected output patterns (CHECK lines)
31+
list<string> checkLines = check;
32+
}
33+
34+
// Base class for elements that can have auto-generated LIT tests
35+
class Testable {
36+
// List of LIT tests associated with this element
37+
list<LitTest> tests = [];
38+
}
39+
40+
#endif // TESTABLE

mlir/include/mlir/Pass/PassBase.td

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
#ifndef MLIR_PASS_PASSBASE
1515
#define MLIR_PASS_PASSBASE
1616

17+
include "mlir/IR/Testable.td"
18+
1719
//===----------------------------------------------------------------------===//
1820
// Options
1921
//===----------------------------------------------------------------------===//
@@ -62,7 +64,7 @@ class Statistic<string varName, string statName, string desc> {
6264
// Pass
6365
//===----------------------------------------------------------------------===//
6466

65-
class PassBase<string passArg, string base> {
67+
class PassBase<string passArg, string base> : Testable {
6668
// The command line argument of the pass.
6769
string argument = passArg;
6870

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// RUN: mlir-tblgen -gen-lit-tests -I %S/../../include -dialect=test %s | FileCheck %s
2+
3+
include "mlir/Pass/PassBase.td"
4+
include "mlir/IR/Testable.td"
5+
6+
def TestPassWithEmbeddedLitTests : Pass<"test-pass-with-embedded-lit-tests"> {
7+
let summary = "pass summary";
8+
let description = [{
9+
Pass description
10+
}];
11+
12+
let tests = [
13+
LitTest<
14+
"lit_test_file_1.mlir",
15+
[{
16+
func.func @test1() {
17+
return 42;
18+
}
19+
}],
20+
[
21+
"// RUN: mlir-opt %s --verify-roundtrip | FileCheck %s",
22+
],
23+
[
24+
"// RANDOM-CHECK-LABEL: func.func @test1",
25+
]
26+
>,
27+
LitTest<
28+
"lit_test_file_2.mlir",
29+
[{
30+
func.func @test2() {
31+
return 42;
32+
}
33+
}],
34+
[
35+
"// RUN: mlir-opt %s --verify-roundtrip | FileCheck %s",
36+
],
37+
[
38+
"// RANDOM-CHECK-LABEL: func.func @test2",
39+
]
40+
>,
41+
];
42+
}
43+
44+
// CHECK-LABEL: // Generated 2 LIT test files
45+
// CHECK: // Use the following files for LIT testing:
46+
47+
// CHECK: // File: generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir
48+
// CHECK: // --- BEGIN generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir ---
49+
// CHECK: // RUN: mlir-opt %s --verify-roundtrip | FileCheck %s
50+
// CHECK: // Generated from TableGen definition: TestPassWithEmbeddedLitTests
51+
// CHECK: func.func @test1() {
52+
// CHECK: return 42;
53+
// CHECK: }
54+
// CHECK: // RANDOM-CHECK-LABEL: func.func @test1
55+
// CHECK: --- END generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir ---
56+
57+
// CHECK: // File: generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir
58+
// CHECK: // --- BEGIN generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir ---
59+
// CHECK: // RUN: mlir-opt %s --verify-roundtrip | FileCheck %s
60+
// CHECK: // Generated from TableGen definition: TestPassWithEmbeddedLitTests
61+
// CHECK: func.func @test2() {
62+
// CHECK: return 42;
63+
// CHECK: }
64+
// CHECK: // RANDOM-CHECK-LABEL: func.func @test2
65+
// CHECK: // --- END generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir ---

mlir/tools/mlir-tblgen/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ add_tablegen(mlir-tblgen MLIR
1616
EnumsGen.cpp
1717
EnumPythonBindingGen.cpp
1818
FormatGen.cpp
19+
LitTestGen.cpp
1920
LLVMIRConversionGen.cpp
2021
LLVMIRIntrinsicGen.cpp
2122
mlir-tblgen.cpp

0 commit comments

Comments
 (0)