Skip to content

[mlir][tblgen] Adds support for embedded LIT tests in TableGen records #158017

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][tblgen] Adds support for embedded LIT tests in TableGen records #158017

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 100 additions & 0 deletions mlir/cmake/modules/AddMLIR.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -762,3 +762,103 @@ function(mlir_target_link_libraries target type)
target_link_libraries(${target} ${type} ${ARGN})
endif()
endfunction()

# Extracts LIT tests embedded in `Testable` records in `tblgen_file`
# and generates a file per test in `output_dir`
#
# Example usage:
# # Extract tests from MyPasses.td and generate them in test/Passes/
# add_embedded_lit_tests(MyPassesEmbeddedTests
# ${CMAKE_CURRENT_SOURCE_DIR}/include/MyPasses.td
# ${CMAKE_CURRENT_SOURCE_DIR}/test/Passes/)
#
# # This will:
# # 1. Process MyPasses.td with mlir-tblgen --gen-lit-tests
# # 2. Extract individual test files to test/Passes/
# # 3. Generate files like: test/Passes/generated_MyPass_test1.mlir
#
function(add_embedded_lit_tests target tblgen_file output_dir)
set(LLVM_TARGET_DEFINITIONS ${tblgen_file})

# Extraction script content
set(EXTRACT_SCRIPT_CONTENT [[
# Generated extraction script
if(NOT CONSOLIDATED_FILE)
message(FATAL_ERROR "CONSOLIDATED_FILE variable is required")
endif()

if(NOT OUTPUT_DIR)
message(FATAL_ERROR "OUTPUT_DIR variable is required")
endif()

if(NOT EXISTS ${CONSOLIDATED_FILE})
message(FATAL_ERROR "Consolidated file does not exist: ${CONSOLIDATED_FILE}")
endif()

# Read the consolidated file
file(READ ${CONSOLIDATED_FILE} file_content)

# Split into lines for processing
string(REPLACE "\n" ";" lines "${file_content}")

set(current_filename "")
set(current_content "")
set(in_test_block FALSE)
set(extracted_test_files)

foreach(line IN LISTS lines)
# Check for filename line
if(line MATCHES "^// File: (.+)$")
set(current_filename "${CMAKE_MATCH_1}")
endif()

# Check for BEGIN marker
if(line MATCHES "^// --- BEGIN .+ ---$")
set(in_test_block TRUE)
set(current_content "")
# Check for END marker
elseif(line MATCHES "^// --- END .+ ---$")
set(in_test_block FALSE)

# Write the extracted content to file
if(current_filename AND current_content)
file(MAKE_DIRECTORY ${OUTPUT_DIR})
file(WRITE ${OUTPUT_DIR}/${current_filename} "${current_content}")
message(STATUS "Extracted test file: ${current_filename}")
list(APPEND extracted_test_files ${current_filename})
endif()

set(current_filename "")
set(current_content "")
# Collect content within BEGIN/END block
elseif(in_test_block)
string(APPEND current_content "${line}\n")
endif()
endforeach()

list(LENGTH extracted_test_files num_extracted_files)
message(STATUS "Extracted ${num_extracted_files} test files to ${OUTPUT_DIR}")
]])

# Write extraction script to a file in the build directory
file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/extract_lit_tests.cmake "${EXTRACT_SCRIPT_CONTENT}")

# Process tblgen_file and generate a file with all embedded LIT
# tests in tblgen_file
get_filename_component(tblgen_name ${tblgen_file} NAME_WE)
set(consolidated_output_file ${tblgen_name}_extracted_lit_tests.txt)
mlir_tablegen(${consolidated_output_file} --gen-lit-tests)

# Add public tablegen target to trigger builds on changes in tblgen_file
add_public_tablegen_target(${target})

# Call the extraction script to extract all LIT tests into individual
# `.mlir` test files
add_custom_command(TARGET ${target} POST_BUILD
COMMAND ${CMAKE_COMMAND}
-DCONSOLIDATED_FILE=${CMAKE_CURRENT_BINARY_DIR}/${consolidated_output_file}
-DOUTPUT_DIR=${output_dir}
-P ${CMAKE_CURRENT_BINARY_DIR}/extract_lit_tests.cmake
COMMENT "Extracting LIT tests to individual files"
)
endfunction()
40 changes: 40 additions & 0 deletions mlir/include/mlir/IR/Testable.td
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
//===-- Testable.td - Testable type definition file --------*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains the definition of the `Testable` type.
//
// Any type whose records can have corresponding LIT tests (eg - Pass) can extend
// `Testable` in order to be able to embed LIT tests within record definitions.
//
//===----------------------------------------------------------------------===//

#ifndef TESTABLE
#define TESTABLE

// Represents a LIT test record in TableGen
class LitTest<string name, code snippet, list<string> run = [], list<string> check = []> {
// The name of the generated test file
string testFileName = name;

// The IR snippet/code to be tested
code irSnippet = snippet;

// The RUN commands for the test (e.g., "mlir-opt %s")
list<string> runLines = run;

// Expected output patterns (CHECK lines)
list<string> checkLines = check;
}

// Base class for elements that can have auto-generated LIT tests
class Testable {
// List of LIT tests associated with this element
list<LitTest> tests = [];
}

#endif // TESTABLE
4 changes: 3 additions & 1 deletion mlir/include/mlir/Pass/PassBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
#ifndef MLIR_PASS_PASSBASE
#define MLIR_PASS_PASSBASE

include "mlir/IR/Testable.td"

//===----------------------------------------------------------------------===//
// Options
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -62,7 +64,7 @@ class Statistic<string varName, string statName, string desc> {
// Pass
//===----------------------------------------------------------------------===//

class PassBase<string passArg, string base> {
class PassBase<string passArg, string base> : Testable {
// The command line argument of the pass.
string argument = passArg;

Expand Down
65 changes: 65 additions & 0 deletions mlir/test/mlir-tblgen/gen-lit-tests.td
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
// RUN: mlir-tblgen -gen-lit-tests -I %S/../../include -dialect=test %s | FileCheck %s

include "mlir/Pass/PassBase.td"
include "mlir/IR/Testable.td"

def TestPassWithEmbeddedLitTests : Pass<"test-pass-with-embedded-lit-tests"> {
let summary = "pass summary";
let description = [{
Pass description
}];

let tests = [
LitTest<
"lit_test_file_1.mlir",
[{
func.func @test1() {
Copy link
Contributor

@j2kun j2kun Sep 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than try to reconstruct the list test from separate records, it seems like it would be simpler to have a single text field that contains the entire lit test, and this feature would limit the allowed complexity of the RUN line. (e.g., only one RUN line with one pipe)

Or you could have the entire lit test in one field, except for the RUN line, and autogenerate the run line as // RUN: mlir-opt %s <pass-flag> | FileCheck %s, since <pass-flag> is already part of the tablegen record and also it would be a good practice to ensure that the behavior tested/documented for a pass is restricted to just what that one pass does.

If you render the before/after example from the lit test, then, you could just strip out any lines starting with \s*// CHECK.

Copy link
Collaborator

@joker-eph joker-eph Sep 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel the same, we actually have a prototype downstream.

  let mlirExamples = [[{
    # func.func @std_for(%arg0 : index, %arg1 : index, %arg2 : index) {
        // Two nested for loops
        scf.for %i0 = %arg0 to %arg1 step %arg2 {
          scf.for %i1 = %arg0 to %arg1 step %arg2 {
            %min_cmp = arith.cmpi slt, %i0, %i1 : index
            %min = arith.select %min_cmp, %i0, %i1 : index
            %max_cmp = arith.cmpi sge, %i0, %i1 : index
            %max = arith.select %max_cmp, %i0, %i1 : index
            // A loop with an empty body
            scf.for %i2 = %min to %max step %i1 {
            }
          }
        }
    # return
   # }
  }]];

We automatically generate the round-trip testing from this, and use it to generate an "example" section in the documentation, while stripping the lines starting with # from the documentation, which allows to build more complex example without clobbering the documentation with the surrounding noise.

Copy link
Collaborator

@joker-eph joker-eph Sep 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of a separate tablegen field, I would consider some markdown in the description field, that we can do:

let description = [[{
  This op is doing ... and ...
  The basic usage is .... 
  ```mlir_example
   # func.func @foo(....) {
      "my.op"() ......
   # }
   ` ``
   It optional takes an extra flag to indicate ...
  ```mlir_example
   # func.func @foo(....) {
      "my.op"() some_flag ......
   # }
   ` ``
   etc.
}]]];

This can then be similarly extracted for round-trip test generation by default, while also integrating nicely when generating the markdown documentation (eliding the commented out lines).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the suggestions. In general I agree with the ideas you guys mentioned. I have a few questions to get a better understanding of the desired end product.

Do we always want to include the IR examples in the documentation or do we want to give the user an option to only generate LIT tests from them?

  1. If it's the former then we'd always need the user to also specify the full "after" version of the IR. This "after" version obviously can be used as the CHECK lines in the test. It does increase the onus on the user a bit, but should give high-quality before/after versions of the IR to go along with documentation, as well as precise LIT tests.

  2. If it's the latter, we'd still need the user to specify CHECK lines, but It'd be quite complicated to generate the "after" versions of the IR for documentation.

  3. It seemed to me, from the comments, that there may be a third option of just generating a --verify-roundtrip test.

    This can then be similarly extracted for round-trip test generation by default...

    If so, I'm not sure what utility such IR can provide either for documentation or for testing.

I personally prefer option 1, and have the user always always specify IR examples in before/after pairs.

Thoughts?

Copy link
Collaborator

@joker-eph joker-eph Sep 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we always want to include the IR examples in the documentation or do we want to give the user an option to only generate LIT tests from them?

If it's not the documentation, I'd question why is it in ODS?

If it's the former then we'd always need the user to also specify the full "after" version of the IR.

I don't quite get what is the "after" version in the context of the documentation?

It seemed to me, from the comments, that there may be a third option of just generating a --verify-roundtrip test.

Right.

If so, I'm not sure what utility such IR can provide either for documentation or for testing.

I don't follow: the IR here isn't meant to be useful for testing, we already have tests for it. The problem statement as I understand it is that we want examples in the doc that are maintained up-to-date in terms of being "correct". That is: we want to test the documentation to actually have examples reflecting the implementation.

Is there another problem to solve that I'm missing?

Copy link
Contributor Author

@jkshtj jkshtj Sep 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it's the former then we'd always need the user to also specify the full "after" version of the IR.

I don't quite get what is the "after" version in the context of the documentation?

It seemed to me, from the comments, that there may be a third option of just generating a --verify-roundtrip test.

Right.

Ok, I see the misalignment in our expectations now.

You're primarily thinking about tests for operations while I started out with working on tests for passes. For operations, I agree that all we need are --verify-roundtrip tests. For passes however, I think it makes sense to have the user specify before/after IR snippets that we generate the tests from.

And I do like the idea of having the snippets be a part of the description field itself, so as to not introduce additional complications for documentation rendering. Let me take another shot at things based on the feedback.

Thank you for your comments!

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh passes is a whole other world, I'm not sure what would make sense for passes, I would need to see any proposal with concrete example on a few actual passes in the code base to figure out.

return 42;
}
}],
[
"// RUN: mlir-opt %s --verify-roundtrip | FileCheck %s",
],
[
"// RANDOM-CHECK-LABEL: func.func @test1",
]
>,
LitTest<
"lit_test_file_2.mlir",
[{
func.func @test2() {
return 42;
}
}],
[
"// RUN: mlir-opt %s --verify-roundtrip | FileCheck %s",
],
[
"// RANDOM-CHECK-LABEL: func.func @test2",
]
>,
];
}

// CHECK-LABEL: // Generated 2 LIT test files
// CHECK: // Use the following files for LIT testing:

// CHECK: // File: generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir
// CHECK: // --- BEGIN generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir ---
// CHECK: // RUN: mlir-opt %s --verify-roundtrip | FileCheck %s
// CHECK: // Generated from TableGen definition: TestPassWithEmbeddedLitTests
// CHECK: func.func @test1() {
// CHECK: return 42;
// CHECK: }
// CHECK: // RANDOM-CHECK-LABEL: func.func @test1
// CHECK: --- END generated_TestPassWithEmbeddedLitTests_lit_test_file_1.mlir ---

// CHECK: // File: generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir
// CHECK: // --- BEGIN generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir ---
// CHECK: // RUN: mlir-opt %s --verify-roundtrip | FileCheck %s
// CHECK: // Generated from TableGen definition: TestPassWithEmbeddedLitTests
// CHECK: func.func @test2() {
// CHECK: return 42;
// CHECK: }
// CHECK: // RANDOM-CHECK-LABEL: func.func @test2
// CHECK: // --- END generated_TestPassWithEmbeddedLitTests_lit_test_file_2.mlir ---
1 change: 1 addition & 0 deletions mlir/tools/mlir-tblgen/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ add_tablegen(mlir-tblgen MLIR
EnumsGen.cpp
EnumPythonBindingGen.cpp
FormatGen.cpp
LitTestGen.cpp
LLVMIRConversionGen.cpp
LLVMIRIntrinsicGen.cpp
mlir-tblgen.cpp
Expand Down
Loading
Loading