[Mlir-commits] [mlir] [mlir][tblgen] Adds support for embedded LIT tests in TableGen records (PR #158017)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Sep 11 01:59:48 PDT 2025


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Kshitij Jain (jkshtj)

<details>
<summary>Changes</summary>

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)
```

---
Full diff: https://github.com/llvm/llvm-project/pull/158017.diff


6 Files Affected:

- (modified) mlir/cmake/modules/AddMLIR.cmake (+100) 
- (added) mlir/include/mlir/IR/Testable.td (+40) 
- (modified) mlir/include/mlir/Pass/PassBase.td (+3-1) 
- (added) mlir/test/mlir-tblgen/gen-lit-tests.td (+65) 
- (modified) mlir/tools/mlir-tblgen/CMakeLists.txt (+1) 
- (added) mlir/tools/mlir-tblgen/LitTestGen.cpp (+170) 


``````````diff
diff --git a/mlir/cmake/modules/AddMLIR.cmake b/mlir/cmake/modules/AddMLIR.cmake
index 6589458ab7894..9b05b70231dba 100644
--- a/mlir/cmake/modules/AddMLIR.cmake
+++ b/mlir/cmake/modules/AddMLIR.cmake
@@ -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()
\ No newline at end of file
diff --git a/mlir/include/mlir/IR/Testable.td b/mlir/include/mlir/IR/Testable.td
new file mode 100644
index 0000000000000..15814ed1bd939
--- /dev/null
+++ b/mlir/include/mlir/IR/Testable.td
@@ -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
\ No newline at end of file
diff --git a/mlir/include/mlir/Pass/PassBase.td b/mlir/include/mlir/Pass/PassBase.td
index e37f9735e2241..50ea44419ca24 100644
--- a/mlir/include/mlir/Pass/PassBase.td
+++ b/mlir/include/mlir/Pass/PassBase.td
@@ -14,6 +14,8 @@
 #ifndef MLIR_PASS_PASSBASE
 #define MLIR_PASS_PASSBASE
 
+include "mlir/IR/Testable.td"
+
 //===----------------------------------------------------------------------===//
 // Options
 //===----------------------------------------------------------------------===//
@@ -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;
 
diff --git a/mlir/test/mlir-tblgen/gen-lit-tests.td b/mlir/test/mlir-tblgen/gen-lit-tests.td
new file mode 100644
index 0000000000000..40a03fb2b2d60
--- /dev/null
+++ b/mlir/test/mlir-tblgen/gen-lit-tests.td
@@ -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() {
+            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 ---
\ No newline at end of file
diff --git a/mlir/tools/mlir-tblgen/CMakeLists.txt b/mlir/tools/mlir-tblgen/CMakeLists.txt
index 2a7ef7e0576c8..e721f1e26a2bd 100644
--- a/mlir/tools/mlir-tblgen/CMakeLists.txt
+++ b/mlir/tools/mlir-tblgen/CMakeLists.txt
@@ -16,6 +16,7 @@ add_tablegen(mlir-tblgen MLIR
   EnumsGen.cpp
   EnumPythonBindingGen.cpp
   FormatGen.cpp
+  LitTestGen.cpp
   LLVMIRConversionGen.cpp
   LLVMIRIntrinsicGen.cpp
   mlir-tblgen.cpp
diff --git a/mlir/tools/mlir-tblgen/LitTestGen.cpp b/mlir/tools/mlir-tblgen/LitTestGen.cpp
new file mode 100644
index 0000000000000..49a092fa9879f
--- /dev/null
+++ b/mlir/tools/mlir-tblgen/LitTestGen.cpp
@@ -0,0 +1,170 @@
+//===- LitTestGen.cpp - LIT test generator ----------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// LitTestGen extracts `LitTest` records from `Testable` TableGen records and 
+// generates corresponding LIT test files.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/TableGen/GenInfo.h"
+#include "mlir/TableGen/Operator.h"
+#include "mlir/TableGen/Pass.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/Path.h"
+#include "llvm/TableGen/Error.h"
+#include "llvm/TableGen/Record.h"
+
+#include <set>
+
+using namespace mlir;
+using namespace mlir::tblgen;
+using llvm::formatv;
+using llvm::RecordKeeper;
+
+static llvm::cl::OptionCategory litTestGenCategory("Options for -gen-lit-tests");
+static llvm::cl::opt<std::string>
+    outputDir("output-dir", 
+              llvm::cl::desc("Output directory for generated test files"),
+              llvm::cl::cat(litTestGenCategory), 
+              llvm::cl::value_desc("directory"));
+
+
+/// Cpp type corresponding to the `LitTest` record type in TableGen
+struct LitTest {
+  std::string sourceDefName;
+  std::string testFileName;
+  std::string irSnippet;  
+  llvm::SmallVector<std::string> runLines;
+  llvm::SmallVector<std::string> checkLines;
+};
+
+static llvm::SmallVector<LitTest> extractTestsFromRecord(const llvm::Record *record,
+                                                         llvm::StringRef dialectName = "") {
+  llvm::SmallVector<LitTest> tests;
+  
+  // Check if the record has a tests field
+  const llvm::RecordVal *testsVal = record->getValue("tests");
+  if (!testsVal)
+    return tests;
+    
+  const llvm::ListInit *testsList = 
+    llvm::dyn_cast_or_null<llvm::ListInit>(testsVal->getValue());
+  if (!testsList)
+    return tests;
+    
+  for (const llvm::Init *init : testsList->getElements()) {
+    const llvm::DefInit *defInit = llvm::dyn_cast<llvm::DefInit>(init);
+    if (!defInit)
+      continue;
+      
+    const llvm::Record *testRec = defInit->getDef();
+    
+    // Extract fields from LitTest record
+    std::string name = testRec->getValueAsString("testFileName").str();
+    std::string irSnippet = testRec->getValueAsString("irSnippet").str();
+    
+    llvm::SmallVector<std::string> runLines;
+    llvm::for_each(*testRec->getValueAsListInit("runLines"), [&](const llvm::Init *init) {
+      runLines.emplace_back(llvm::cast<llvm::StringInit>(init)->getValue());
+    });
+
+    llvm::SmallVector<std::string> checkLines;
+    llvm::for_each(*testRec->getValueAsListInit("checkLines"), [&](const llvm::Init *init) {
+      checkLines.emplace_back(llvm::cast<llvm::StringInit>(init)->getValue());
+    });
+
+    tests.push_back(LitTest {
+      record->getName().str(),
+      name, 
+      irSnippet, 
+      runLines, 
+      checkLines, 
+    });
+  }
+  
+  return tests;
+}
+
+/// Extract tests from passes
+static llvm::SmallVector<LitTest> extractPassTests(const RecordKeeper &records) {
+  llvm::SmallVector<LitTest> tests;
+  
+  // Check if PassBase class exists before trying to get derived definitions
+  if (records.getClass("PassBase")) {
+    for (const llvm::Record *def : records.getAllDerivedDefinitions("PassBase")) {
+      if (def->isAnonymous())
+        continue;
+        
+      auto passTests = extractTestsFromRecord(def, "passes");
+      tests.insert(tests.end(), passTests.begin(), passTests.end());
+    }
+  }
+  
+  return tests;
+}
+
+/// Generate a LIT test file for an IR test
+static void generateTestFile(const LitTest &test, llvm::raw_ostream &os) {
+  // Add RUN lines
+  for (const auto& runLine : test.runLines) {
+    os << "\n" << runLine << "\n";
+  }
+
+  os << "// Generated from TableGen definition: " << test.sourceDefName << "\n\n";
+  
+  // Add the test body
+  os << test.irSnippet << "\n";
+  
+  // Add CHECK lines
+  for (const auto& checkLine : test.checkLines) {
+    os << "\n" << checkLine << "\n";
+  }
+}
+
+/// Main function to generate all IR test test files
+static void generateLitTests(const RecordKeeper &records, raw_ostream &os) {
+  llvm::SmallVector<LitTest> allTests;
+  
+  // Extract tests from different definition types (only passes for now)
+  auto passTests = extractPassTests(records);
+  
+  allTests.insert(allTests.end(), passTests.begin(), passTests.end());
+  
+  if (allTests.empty()) {
+    os << "// No LitTest record found in any TableGen definition\n";
+    return;
+  }
+  
+  // Generate summary
+  os << "// Generated " << allTests.size() << " LIT test files\n";
+  os << "// Use the following files for LIT testing:\n\n";
+  
+  // Generate file list and content for each test
+  for (const auto& test : allTests) {
+    std::string testFileName = formatv("generated_{0}_{1}", test.sourceDefName, test.testFileName);
+    os << "// File: " << testFileName << "\n";
+    
+    os << "// --- BEGIN " << testFileName << " ---\n";
+    generateTestFile(test, os);
+    os << "// --- END " << testFileName << " ---\n\n";
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Generator Registration
+//===----------------------------------------------------------------------===//
+
+static mlir::GenRegistration
+    genLitTests("gen-lit-tests", "Generate LIT test files for `Testable` TableGen records",
+                  [](const RecordKeeper &records, raw_ostream &os) {
+                    generateLitTests(records, os);
+                    return false;
+                  });
\ No newline at end of file

``````````

</details>


https://github.com/llvm/llvm-project/pull/158017


More information about the Mlir-commits mailing list