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

Kshitij Jain llvmlistbot at llvm.org
Sun Nov 9 12:52:17 PST 2025


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

>From 73613958573a098e4b3785bad0d25528404b6164 Mon Sep 17 00:00:00 2001
From: Kshitij <kjain at d-matrix.ai>
Date: Thu, 11 Sep 2025 08:32:03 +0000
Subject: [PATCH 1/2] [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)
---
 mlir/cmake/modules/AddMLIR.cmake       | 100 +++++++++++++++
 mlir/include/mlir/IR/Testable.td       |  40 ++++++
 mlir/include/mlir/Pass/PassBase.td     |   4 +-
 mlir/test/mlir-tblgen/gen-lit-tests.td |  65 ++++++++++
 mlir/tools/mlir-tblgen/CMakeLists.txt  |   1 +
 mlir/tools/mlir-tblgen/LitTestGen.cpp  | 170 +++++++++++++++++++++++++
 6 files changed, 379 insertions(+), 1 deletion(-)
 create mode 100644 mlir/include/mlir/IR/Testable.td
 create mode 100644 mlir/test/mlir-tblgen/gen-lit-tests.td
 create mode 100644 mlir/tools/mlir-tblgen/LitTestGen.cpp

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

>From d0005879c4e983b4f7726e53e60ed3e55f5286c5 Mon Sep 17 00:00:00 2001
From: jkshtj <jkshtj at outlook.com>
Date: Sun, 9 Nov 2025 12:43:01 -0800
Subject: [PATCH 2/2] R1 - Addresses feedback

---
 mlir/cmake/modules/AddMLIR.cmake         |  12 +-
 mlir/docs/DefiningDialects/Operations.md |  87 +++++++++++
 mlir/include/mlir/IR/Testable.td         |  40 -----
 mlir/include/mlir/Pass/PassBase.td       |   4 +-
 mlir/test/mlir-tblgen/gen-lit-tests.td   | 123 ++++++++-------
 mlir/tools/mlir-tblgen/LitTestGen.cpp    | 187 ++++++++++++++---------
 mlir/tools/mlir-tblgen/OpDocGen.cpp      |  62 +++++++-
 7 files changed, 331 insertions(+), 184 deletions(-)
 delete mode 100644 mlir/include/mlir/IR/Testable.td

diff --git a/mlir/cmake/modules/AddMLIR.cmake b/mlir/cmake/modules/AddMLIR.cmake
index 9b05b70231dba..de61303bbdcbd 100644
--- a/mlir/cmake/modules/AddMLIR.cmake
+++ b/mlir/cmake/modules/AddMLIR.cmake
@@ -763,19 +763,19 @@ function(mlir_target_link_libraries target type)
   endif()
 endfunction()
 
-# Extracts LIT tests embedded in `Testable` records in `tblgen_file`
+# Extracts LIT tests embedded in op descriptions 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/)
+#                          ${CMAKE_CURRENT_SOURCE_DIR}/include/MyDialectOps.td
+#                          ${CMAKE_CURRENT_SOURCE_DIR}/test/Dialect/MyDialect)
 #
 #   # 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
+#   # 1. Process MyDialectOps.td with mlir-tblgen --gen-lit-tests
+#   # 2. Extract individual test files to test/Dialect/MyDialect
+#   # 3. Generate files like: test/Dialect/MyDialect/generated_MyOp_1.mlir
 #
 function(add_embedded_lit_tests target tblgen_file output_dir)
   set(LLVM_TARGET_DEFINITIONS ${tblgen_file})
diff --git a/mlir/docs/DefiningDialects/Operations.md b/mlir/docs/DefiningDialects/Operations.md
index f988bebea1223..b1e0b0baa7934 100644
--- a/mlir/docs/DefiningDialects/Operations.md
+++ b/mlir/docs/DefiningDialects/Operations.md
@@ -175,6 +175,93 @@ understanding the operation.
 >     starting with a capital letter and without trailing punctuation.
 >     Put expanded explanation in the description.
 
+#### Autogenerated LIT tests from IR examples in operation description
+
+To help ensure that IR based op examples stay synchronized with
+the actual operation definitions and remain valid as the codebase evolves,
+MLIR supports automatically generating roundtrip verification LIT tests from
+code examples embedded in operation descriptions.
+
+##### Embedding an IR example in op description
+
+To embed a testable MLIR example in an operation's description, use a special
+markdown code block with the `mlir_example` tag:
+
+```tablegen
+def MyOp : MyDialect_Op<"my_op"> {
+  let summary = "An example operation";
+
+  let description = [{
+    This operation demonstrates example usage.
+
+    ```mlir_example
+    # func.func @example(%arg0: i32) -> i32 { // This line will be ignored when generating `.md` documentation
+      %result = mydialect.my_op %arg0 : i32
+      return %result : i32
+    # } // This line is ignored as well
+    ```
+
+     ```mlir_example(custom-driver)
+    # func.func @example(%arg0: i32) -> i32 {
+      %result = mydialect.my_op %arg0 : i32
+      return %result : i32
+    # }
+    ```
+  }];
+}
+```
+
+Things to note:
+
+1. Any portion of the embedded snippet that need not be included in the
+   `.md` documentation can be demarcated with a `#` prefix. Lines starting with
+   `#` are stripped out when generating operation documentation but are included
+   (without the `#` and indentation) in the generated LIT tests. This allows you
+   to provide context (like function wrappers) needed for valid MLIR IR without
+   cluttering the documentation.
+
+2. You can optionally specify a custom tool/driver to use in the generated RUN
+   line by adding it in parentheses after `mlir_example`, such as
+   `mlir_example(custom-driver)`. If not specified, `mlir-opt` is used by default.
+
+3. Multiple `mlir_example` blocks can be included in a single operation's
+   description. Each block will generate a separate test file with a unique name
+   based on the operation name and example index (e.g.,
+   `generated_MyOp_example_0.mlir`, `generated_MyOp_example_1.mlir`).
+
+4. The generated RUN line automatically includes the `--verify-roundtrip` flag
+   to ensure the operation can be parsed and printed correctly.
+
+##### CMake integration for LIT test file generation
+
+MLIR provides a CMake function `add_embedded_lit_tests` to automatically extract
+and generate LIT test files from TableGen definitions during the build process.
+This function is defined in `mlir/cmake/modules/AddMLIR.cmake`.
+
+**Usage:**
+
+```cmake
+add_embedded_lit_tests(<target_name> <tablegen_file> <output_directory>)
+```
+
+**Parameters:**
+- `target_name`: Name for the CMake target that will be created
+- `tablegen_file`: Path to the TableGen file containing operations with `mlir_example` blocks
+- `output_directory`: Directory where the extracted test files will be written
+
+**Example:**
+
+```cmake
+# Extract tests from MyDialectOps.td and generate them in test/Dialect/MyDialect/
+add_embedded_lit_tests(MyDialectEmbeddedTests
+                       ${CMAKE_CURRENT_SOURCE_DIR}/include/MyDialect/MyDialectOps.td
+                       ${CMAKE_CURRENT_SOURCE_DIR}/test/Dialect/MyDialect/)
+```
+
+The generated test files can then be picked up by your LIT test configuration
+automatically, ensuring that all embedded examples are continuously tested as
+part of your regular test suite.
+
 ### Operation arguments
 
 There are three kinds of arguments: operands, attributes, and properties.
diff --git a/mlir/include/mlir/IR/Testable.td b/mlir/include/mlir/IR/Testable.td
deleted file mode 100644
index 15814ed1bd939..0000000000000
--- a/mlir/include/mlir/IR/Testable.td
+++ /dev/null
@@ -1,40 +0,0 @@
-//===-- 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 50ea44419ca24..e37f9735e2241 100644
--- a/mlir/include/mlir/Pass/PassBase.td
+++ b/mlir/include/mlir/Pass/PassBase.td
@@ -14,8 +14,6 @@
 #ifndef MLIR_PASS_PASSBASE
 #define MLIR_PASS_PASSBASE
 
-include "mlir/IR/Testable.td"
-
 //===----------------------------------------------------------------------===//
 // Options
 //===----------------------------------------------------------------------===//
@@ -64,7 +62,7 @@ class Statistic<string varName, string statName, string desc> {
 // Pass
 //===----------------------------------------------------------------------===//
 
-class PassBase<string passArg, string base> : Testable {
+class PassBase<string passArg, string base> {
   // 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
index 40a03fb2b2d60..2c57bc9e108dc 100644
--- a/mlir/test/mlir-tblgen/gen-lit-tests.td
+++ b/mlir/test/mlir-tblgen/gen-lit-tests.td
@@ -1,65 +1,74 @@
-// RUN: mlir-tblgen -gen-lit-tests -I %S/../../include -dialect=test %s | FileCheck %s
+// RUN: mlir-tblgen -gen-lit-tests -I %S/../../include %s | FileCheck %s
 
-include "mlir/Pass/PassBase.td"
-include "mlir/IR/Testable.td"
+include "mlir/IR/OpBase.td"
 
-def TestPassWithEmbeddedLitTests : Pass<"test-pass-with-embedded-lit-tests"> {
-  let summary = "pass summary";
+def Test_Dialect : Dialect {
+  let name = "test";
+  let cppNamespace = "test";
+}
+
+def TestOp : Op<Test_Dialect, "test_op"> {
+  let summary = "test op with mlir_example code blocks";
   let description = [{
-    Pass description
+    This operation demonstrates the mlir_example feature for ops.
+
+    Basic usage:
+    ```mlir_example(mlir-opt)
+    # func.func @foo(%arg0: i32) -> i32 {
+      %0 = test.test_op %arg0 : i32
+      return %0 : i32
+    # }
+    ```
+
+    And some more back to back examples -
+
+    ```mlir_example(some-other-tool)
+    # func.func @foo1(%arg1: i32) -> i32 {
+      %0 = test.test_op %arg1 : i32
+      return %0 : i32
+    # }
+    ```
+    ```mlir_example(yet-another-tool)
+    # func.func @foo2(%arg2: i32) -> i32 {
+      %0 = test.test_op %arg2 : i32
+      return %0 : i32
+    # }
+    ```
   }];
-  
-  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",
-      ]
-    >,
-  ];
+
+  let arguments = (ins I32:$input);
+  let results = (outs I32:$output);
 }
 
-// CHECK-LABEL:       // Generated 2 LIT test files
-// CHECK:             // Use the following files for LIT testing:
+// CHECK-LABEL:     // Generated 3 LIT test files
+// CHECK:           // Use the following files for LIT testing:
+
+// CHECK:           // File: generated_TestOp_example_0.mlir
+// CHECK:           // --- BEGIN generated_TestOp_example_0.mlir ---
+// CHECK:           mlir-opt %s --verify-roundtrip
+// CHECK:           // Generated from TableGen definition: TestOp
+// CHECK:           func.func @foo(%arg0: i32) -> i32 {
+// CHECK:           %0 = test.test_op %arg0 : i32
+// CHECK:           return %0 : i32
+// CHECK:           }
+// CHECK:           // --- END generated_TestOp_example_0.mlir ---
 
-// 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_TestOp_example_1.mlir
+// CHECK:           // --- BEGIN generated_TestOp_example_1.mlir ---
+// CHECK:           some-other-tool %s --verify-roundtrip
+// CHECK:           // Generated from TableGen definition: TestOp
+// CHECK:           func.func @foo1(%arg1: i32) -> i32 {
+// CHECK:           %0 = test.test_op %arg1 : i32
+// CHECK:           return %0 : i32
+// CHECK:           }
+// CHECK:           // --- END generated_TestOp_example_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
+// CHECK:           // File: generated_TestOp_example_2.mlir
+// CHECK:           // --- BEGIN generated_TestOp_example_2.mlir ---
+// CHECK:           yet-another-tool %s --verify-roundtrip
+// CHECK:           // Generated from TableGen definition: TestOp
+// CHECK:           func.func @foo2(%arg2: i32) -> i32 {
+// CHECK:           %0 = test.test_op %arg2 : i32
+// CHECK:           return %0 : i32
+// CHECK:           }
+// CHECK:           // --- END generated_TestOp_example_2.mlir ---
\ No newline at end of file
diff --git a/mlir/tools/mlir-tblgen/LitTestGen.cpp b/mlir/tools/mlir-tblgen/LitTestGen.cpp
index 49a092fa9879f..09c31dac063ca 100644
--- a/mlir/tools/mlir-tblgen/LitTestGen.cpp
+++ b/mlir/tools/mlir-tblgen/LitTestGen.cpp
@@ -1,4 +1,5 @@
-//===- LitTestGen.cpp - LIT test generator ----------------------------------===//
+//===- 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.
@@ -6,8 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// LitTestGen extracts `LitTest` records from `Testable` TableGen records and 
-// generates corresponding LIT test files.
+// Extracts embedded LIT tests from operation descriptions in tablegen.
 //
 //===----------------------------------------------------------------------===//
 
@@ -19,11 +19,10 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/Regex.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
-#include <set>
-
 using namespace mlir;
 using namespace mlir::tblgen;
 using llvm::formatv;
@@ -36,82 +35,114 @@ static llvm::cl::opt<std::string>
               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;
+  std::string irSnippet;
+  llvm::SmallVector<std::string, 1> runLines;
   llvm::SmallVector<std::string> checkLines;
 };
 
-static llvm::SmallVector<LitTest> extractTestsFromRecord(const llvm::Record *record,
-                                                         llvm::StringRef dialectName = "") {
+/// Extracts code snippets with mlir_example tag from a description field.
+///
+/// Returns a vector of LitTest objects found within ```mlir_example ... ```
+/// blocks.
+static llvm::SmallVector<LitTest>
+extractOpTests(llvm::StringRef description, llvm::StringRef sourceDefName) {
   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, 
-    });
+
+  // Pattern to match ```mlir_example ... ``` code blocks
+  // - ``` - Three literal backticks
+  // - `mlir_example` - Literal text
+  // - `(\(.+\))?` - Capture group matching the optional RUN tool name. Default
+  // is `mlir-opt`.
+  // - `([^`|`^`|``^`]+)` - Capture group matching the actual mlir IR example
+  // content (everything except for three consecutive backticks).
+  // - ``` - Three literal closing backticks
+  llvm::Regex codeBlockRegex(
+      "```mlir_example(\\([[:alnum:]_-]+\\))?[[:space:]]([^`|`^`|``^`]+)```");
+
+  auto remaining = description;
+  llvm::SmallVector<llvm::StringRef> matches;
+
+  while (codeBlockRegex.match(remaining, &matches)) {
+    if (matches.size() == 3) {
+      std::string tool = "mlir-opt";
+      // matches[1] contains the RUN tool name
+      if (!matches[1].empty()) {
+        tool = matches[1].ltrim('(').rtrim(')').str();
+      }
+
+      // matches[2] contains the code content
+      auto codeRef = matches[2];
+      // Remove leading/trailing whitespace and comment markers (# prefix)
+      llvm::SmallVector<llvm::StringRef> lines;
+      codeRef.split(lines, '\n', -1, false);
+
+      std::string processedCode;
+      for (llvm::StringRef line : lines) {
+        auto isBody = true;
+        line = line.ltrim();
+        // Remove leading # comment markers if present
+        if (line.starts_with("#")) {
+          isBody = false;
+          line = line.drop_front(1).ltrim();
+        }
+        if (!line.empty() || !processedCode.empty()) {
+          auto tab = isBody ? "  " : "";
+          processedCode += tab + line.str() + "\n";
+        }
+      }
+
+      if (!processedCode.empty()) {
+        // Generate test file name based on index
+        auto testFileName = formatv("example_{0}.mlir", tests.size());
+        // Generate default RUN line with --verify-roundtrip
+        auto runLine =
+            llvm::formatv("// RUN: {0} %s --verify-roundtrip", tool).str();
+
+        tests.push_back(LitTest{
+            sourceDefName.str(),
+            testFileName,
+            processedCode,
+            {runLine},
+            {} // No CHECK lines by default
+        });
+      }
+    }
+
+    // Move past this match to find the next one
+    size_t matchEnd =
+        remaining.find("```", remaining.find("```mlir_example") + 15);
+    if (matchEnd == llvm::StringRef::npos)
+      break;
+    remaining = remaining.substr(matchEnd + 3);
   }
-  
+
   return tests;
 }
 
-/// Extract tests from passes
-static llvm::SmallVector<LitTest> extractPassTests(const RecordKeeper &records) {
+static llvm::SmallVector<LitTest>
+extractTestsFromRecord(const llvm::Record *record) {
   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());
-    }
+
+  // Try to extract mlir_example code blocks from the description field
+  const llvm::RecordVal *descVal = record->getValue("description");
+  if (!descVal)
+    return tests;
+
+  auto description = record->getValueAsString("description");
+  if (description.empty())
+    return tests;
+
+  if (record->isSubClassOf("Op")) {
+    tests = extractOpTests(description, record->getName());
   }
-  
+
   return tests;
 }
 
-/// Generate a LIT test file for an IR test
+/// Generates 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) {
@@ -131,27 +162,33 @@ static void generateTestFile(const LitTest &test, llvm::raw_ostream &os) {
 
 /// Main function to generate all IR test test files
 static void generateLitTests(const RecordKeeper &records, raw_ostream &os) {
+  assert(records.getClass("Op") && "Undefined TableGen class type: Op");
+
   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());
-  
+
+  llvm::SmallVector<StringRef, 2> testTypes{"Op"};
+  for (const llvm::Record *def : records.getAllDerivedDefinitions(testTypes)) {
+    if (def->isAnonymous())
+      continue;
+
+    auto opTests = extractTestsFromRecord(def);
+    allTests.insert(allTests.end(), opTests.begin(), opTests.end());
+  }
+
   if (allTests.empty()) {
-    os << "// No LitTest record found in any TableGen definition\n";
+    os << "// No mlir_example code blocks 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";
diff --git a/mlir/tools/mlir-tblgen/OpDocGen.cpp b/mlir/tools/mlir-tblgen/OpDocGen.cpp
index f2b269e3a4542..b89e93588396a 100644
--- a/mlir/tools/mlir-tblgen/OpDocGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDocGen.cpp
@@ -194,6 +194,57 @@ static StringRef resolveAttrDescription(const Attribute &attr) {
   return description;
 }
 
+/// Preprocesses op description.
+///
+/// 1. `mlir_example` code block tag is replaced with `mlir`.
+///
+/// 2. Lines starting with commentMarker inside code blocks (between
+/// snippetStart and snippetEnd) are completely removed.
+static std::string preprocessOpDescription(StringRef description,
+                                           StringRef snippetStart,
+                                           StringRef snippetEnd,
+                                           StringRef commentMarker) {
+  std::string result;
+  raw_string_ostream os(result);
+
+  SmallVector<StringRef> lines;
+  description.split(lines, '\n', -1, false);
+
+  bool insideSnippet = false;
+
+  for (StringRef line : lines) {
+    StringRef trimmedLine = line.ltrim();
+
+    // Check if we're entering a code snippet block
+    if (trimmedLine.starts_with(snippetStart)) {
+      insideSnippet = true;
+      os << "```mlir" << "\n";
+      continue;
+    }
+
+    // Check if we're exiting a code snippet block
+    if (trimmedLine.starts_with(snippetEnd) && insideSnippet) {
+      insideSnippet = false;
+      os << line << "\n";
+      continue;
+    }
+
+    // If we're inside a snippet and the line starts with the comment marker,
+    // skip it
+    if (insideSnippet) {
+      if (trimmedLine.starts_with(commentMarker))
+        continue;
+      else
+        line = line.ltrim();
+    }
+
+    // Otherwise, keep the line as-is
+    os << line << "\n";
+  }
+
+  return result;
+}
+
 static void emitOpDoc(const Operator &op, raw_ostream &os) {
   std::string classNameStr = op.getQualCppClassName();
   StringRef className = classNameStr;
@@ -206,8 +257,11 @@ static void emitOpDoc(const Operator &op, raw_ostream &os) {
   if (op.hasAssemblyFormat())
     emitAssemblyFormat(op.getOperationName(), op.getAssemblyFormat().trim(),
                        os);
-  if (op.hasDescription())
-    mlir::tblgen::emitDescription(op.getDescription(), os);
+  if (op.hasDescription()) {
+    auto preprocessedDescription = preprocessOpDescription(
+        op.getDescription(), "```mlir_example", "```", "#");
+    mlir::tblgen::emitDescription(preprocessedDescription, os);
+  }
 
   emitOpTraitsDoc(op, os);
 
@@ -453,7 +507,9 @@ static void emitBlock(ArrayRef<Attribute> attributes, StringRef inputFilename,
           [&](raw_ostream &os) {
             if (nested) {
               os << "\n## " << StringRef(grouping.summary).trim() << "\n";
-              emitDescription(grouping.description, os);
+              auto preprocessedDescription = preprocessOpDescription(
+                  grouping.description, "```mlir_example", "```", "#");
+              emitDescription(preprocessedDescription, os);
               os << "\n";
             }
             for (const Operator &op : grouping.ops) {



More information about the Mlir-commits mailing list