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

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

mlir/examples/toy/Ch2/CMakeLists.txt

@@ -1,6 +1,11 @@
 # For a better template to copy, see examples/standalone
 add_subdirectory(include)
 
+add_embedded_lit_tests(EmbeddedLitTestsGen 
+  "include/toy/Ops.td"
+  "${CMAKE_CURRENT_BINARY_DIR}"
+)
+
 set(LLVM_LINK_COMPONENTS
   Support
   )
@@ -13,8 +18,9 @@ add_toy_chapter(toyc-ch2
 
   DEPENDS
   ToyCh2OpsIncGen
-
+  EmbeddedLitTestsGen
   )
+
 include_directories(include/)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include/)
 target_link_libraries(toyc-ch2

mlir/include/mlir/Pass/PassBase.td

@@ -14,6 +14,8 @@
 #ifndef MLIR_PASS_PASSBASE
 #define MLIR_PASS_PASSBASE
 
+include "mlir/IR/Testable.td"
+
 //===----------------------------------------------------------------------===//
 // Options
 //===----------------------------------------------------------------------===//

mlir/test/mlir-tblgen/gen-lit-tests.td

@@ -0,0 +1,74 @@
+// RUN: mlir-tblgen -gen-lit-tests -I %S/../../include %s | FileCheck %s
+
+include "mlir/IR/OpBase.td"
+
+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 = [{
+    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 arguments = (ins I32:$input);
+  let results = (outs I32:$output);
+}
+
+// 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_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_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 ---

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