[mlir][PDLL] Add support for generating PDL patterns from PDLL at build time

This essentially sets up mlir-pdll to function in a similar manner to mlir-tblgen. Aside
from the boilerplate of configuring CMake and setting up a basic initial test, two new
options are added to mlir-pdll to mirror options provided by tblgen:

* -d
 This option generates a dependency file (i.e. a set of build time dependencies) while
 processing the input file.

* --write-if-changed
 This option only writes to the output file if the data would have changed, which for
 the build system prevents unnecesarry rebuilds if the file was touched but not actually
 changed.

Differential Revision: https://reviews.llvm.org/D124075

mlir/CMakeLists.txt

@@ -149,6 +149,7 @@ include_directories( ${MLIR_INCLUDE_DIR})
 # from another directory like tools
 add_subdirectory(tools/mlir-tblgen)
 add_subdirectory(tools/mlir-linalg-ods-gen)
+add_subdirectory(tools/mlir-pdll)
 
 add_subdirectory(include/mlir)
 add_subdirectory(lib)

mlir/cmake/modules/AddMLIR.cmake

@@ -7,6 +7,17 @@ function(mlir_tablegen ofn)
       PARENT_SCOPE)
 endfunction()
 
+# Declare a PDLL library in the current directory.
+function(add_mlir_pdll_library target inputFile ofn)
+  set(LLVM_TARGET_DEFINITIONS ${inputFile})
+
+  tablegen(MLIR_PDLL ${ofn} -x=cpp ${ARGN})
+  set(TABLEGEN_OUTPUT ${TABLEGEN_OUTPUT} ${CMAKE_CURRENT_BINARY_DIR}/${ofn}
+      PARENT_SCOPE)
+
+  add_public_tablegen_target(${target})
+endfunction()
+
 # Declare a dialect in the include directory
 function(add_mlir_dialect dialect dialect_namespace)
   set(LLVM_TARGET_DEFINITIONS ${dialect}.td)

mlir/cmake/modules/CMakeLists.txt

@@ -30,6 +30,7 @@ set(MLIR_CONFIG_INCLUDE_DIRS
   )
 # Refer to the best host mlir-tbgen, which might be a host-optimized version
 set(MLIR_CONFIG_TABLEGEN_EXE "${MLIR_TABLEGEN_EXE}")
+set(MLIR_CONFIG_PDLL_TABLEGEN_EXE "${MLIR_PDLL_TABLEGEN_EXE}")
 
 configure_file(
   ${CMAKE_CURRENT_SOURCE_DIR}/MLIRConfig.cmake.in
@@ -60,6 +61,7 @@ set(MLIR_CONFIG_INCLUDE_DIRS
 # Ensure that we are using the installed mlir-tblgen. This might not be MLIR_TABLEGEN_EXE
 # if we're building with a host-optimized mlir-tblgen (with LLVM_OPTIMIZED_TABLEGEN).
 set(MLIR_CONFIG_TABLEGEN_EXE mlir-tblgen)
+set(MLIR_CONFIG_PDLL_TABLEGEN_EXE mlir-pdll)
 
 configure_file(
   ${CMAKE_CURRENT_SOURCE_DIR}/MLIRConfig.cmake.in

mlir/cmake/modules/MLIRConfig.cmake.in

@@ -9,6 +9,7 @@ set(MLIR_EXPORTED_TARGETS "@MLIR_EXPORTS@")
 set(MLIR_CMAKE_DIR "@MLIR_CONFIG_CMAKE_DIR@")
 set(MLIR_INCLUDE_DIRS "@MLIR_CONFIG_INCLUDE_DIRS@")
 set(MLIR_TABLEGEN_EXE "@MLIR_CONFIG_TABLEGEN_EXE@")
+set(MLIR_PDLL_TABLEGEN_EXE "@MLIR_CONFIG_PDLL_TABLEGEN_EXE@")
 set(MLIR_INSTALL_AGGREGATE_OBJECTS "@MLIR_INSTALL_AGGREGATE_OBJECTS@")
 set(MLIR_ENABLE_BINDINGS_PYTHON "@MLIR_ENABLE_BINDINGS_PYTHON@")
 

mlir/test/lib/CMakeLists.txt

@@ -5,4 +5,5 @@ add_subdirectory(IR)
 add_subdirectory(Pass)
 add_subdirectory(Reducer)
 add_subdirectory(Rewrite)
+add_subdirectory(Tools)
 add_subdirectory(Transforms)

mlir/test/lib/Tools/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(PDLL)

mlir/test/lib/Tools/PDLL/CMakeLists.txt

@@ -0,0 +1,34 @@
+add_mlir_pdll_library(MLIRTestPDLLPatternsIncGen
+  TestPDLL.pdll
+  TestPDLLPatterns.h.inc
+
+  EXTRA_INCLUDES
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../Dialect/Test
+  ${CMAKE_CURRENT_BINARY_DIR}/../../Dialect/Test
+  )
+
+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRTestPDLL
+  TestPDLL.cpp
+
+  EXCLUDE_FROM_LIBMLIR
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/PDLL
+
+  DEPENDS
+  MLIRTestPDLLPatternsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRSupport
+  MLIRTestDialect
+  MLIRTransformUtils
+  )
+
+target_include_directories(MLIRTestPDLL
+  PRIVATE
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../Dialect/Test
+  ${CMAKE_CURRENT_BINARY_DIR}/../../Dialect/Test
+  )

mlir/test/lib/Tools/PDLL/TestPDLL.cpp

@@ -0,0 +1,51 @@
+//===- TestPDLByteCode.cpp - Test PDLL functionality ----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+#include "TestPDLLPatterns.h.inc"
+
+namespace {
+struct TestPDLLPass : public PassWrapper<TestPDLLPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPDLLPass)
+
+  StringRef getArgument() const final { return "test-pdll-pass"; }
+  StringRef getDescription() const final { return "Test PDLL functionality"; }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<pdl::PDLDialect, pdl_interp::PDLInterpDialect>();
+  }
+  LogicalResult initialize(MLIRContext *ctx) override {
+    // Build the pattern set within the `initialize` to avoid recompiling PDL
+    // patterns during each `runOnOperation` invocation.
+    RewritePatternSet patternList(&getContext());
+    populateGeneratedPDLLPatterns(patternList);
+    patterns = std::move(patternList);
+    return success();
+  }
+
+  void runOnOperation() final {
+    // Invoke the pattern driver with the provided patterns.
+    (void)applyPatternsAndFoldGreedily(getOperation(), patterns);
+  }
+
+  FrozenRewritePatternSet patterns;
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestPDLLPasses() { PassRegistration<TestPDLLPass>(); }
+} // namespace test
+} // namespace mlir

mlir/test/lib/Tools/PDLL/TestPDLL.pdll

@@ -0,0 +1,12 @@
+//===- TestPDLL.pdll - Test PDLL functionality ----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestOps.td"
+
+/// A simple pattern that matches and replaces an operation.
+Pattern TestSimplePattern => replace op<test.simple> with op<test.success>;

mlir/test/lib/Tools/PDLL/lit.local.cfg

@@ -0,0 +1 @@
+config.suffixes.remove('.pdll')

mlir/test/mlir-pdll/Integration/test-pdll.mlir

@@ -0,0 +1,8 @@
+// RUN: mlir-opt %s -test-pdll-pass | FileCheck %s
+
+// CHECK-LABEL: func @simpleTest
+func @simpleTest() {
+  // CHECK: test.success
+  "test.simple"() : () -> ()
+  return
+}

mlir/test/mlir-pdll/Parser/dependency-file.pdll

@@ -0,0 +1,15 @@
+// RUN: mlir-pdll %s -I %S -I %S/../../../include -d=%t -o %t.cpp.inc
+// RUN: FileCheck %s < %t
+
+// Test support for generating dependency files.
+
+#include "include/ops.td"
+#include "include/included.pdll"
+
+// Check that we depend on the included files. We don't check for all transitive includes
+// here to avoid the need to update this test every time we add a new transitive include.
+// This test is mostly aimed to ensure we are generating the dependency file correctly.
+
+// CHECK: {{.*}}.cpp.inc:
+// CHECK-SAME: include/included.pdll
+// CHECK-SAME: include/ops.td

mlir/tools/CMakeLists.txt

@@ -1,7 +1,6 @@
 add_subdirectory(mlir-lsp-server)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-parser-fuzzer)
-add_subdirectory(mlir-pdll)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-shlib)

mlir/tools/mlir-opt/CMakeLists.txt

@@ -28,6 +28,7 @@ if(MLIR_INCLUDE_TESTS)
     MLIRTestDialect
     MLIRTestIR
     MLIRTestPass
+    MLIRTestPDLL
     MLIRTestReducer
     MLIRTestRewrite
     MLIRTestTransformDialect

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -105,6 +105,7 @@ void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
+void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
@@ -200,6 +201,7 @@ void registerTestPasses() {
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
+  mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestSliceAnalysisPass();

mlir/tools/mlir-pdll/CMakeLists.txt

@@ -1,17 +1,19 @@
-set(LIBS
+set(LLVM_LINK_COMPONENTS
+  Demangle
+  Support
+  TableGen
+)
+
+add_tablegen(mlir-pdll MLIR_PDLL
+  mlir-pdll.cpp
+  )
+
+set_target_properties(mlir-pdll PROPERTIES FOLDER "Tablegenning")
+target_link_libraries(mlir-pdll
+  PRIVATE
   MLIRPDLLAST
   MLIRPDLLCodeGen
   MLIRPDLLParser
   )
 
-add_llvm_tool(mlir-pdll
-  mlir-pdll.cpp
-
-  DEPENDS
-  ${LIBS}
-  )
-
-target_link_libraries(mlir-pdll PRIVATE ${LIBS})
-llvm_update_compile_flags(mlir-pdll)
-
 mlir_check_all_link_libraries(mlir-pdll)

mlir/tools/mlir-pdll/mlir-pdll.cpp

@@ -19,6 +19,7 @@
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include <set>
 
 using namespace mlir;
 using namespace mlir::pdll;
@@ -37,7 +38,7 @@ enum class OutputType {
 static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
               OutputType outputType, std::vector<std::string> &includeDirs,
-              bool dumpODS) {
+              bool dumpODS, std::set<std::string> *includedFiles) {
   llvm::SourceMgr sourceMgr;
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());
@@ -48,6 +49,14 @@ processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
   if (failed(module))
     return failure();
 
+  // Add the files that were included to the set.
+  if (includedFiles) {
+    for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {
+      includedFiles->insert(
+          sourceMgr.getMemoryBuffer(i + 1)->getBufferIdentifier().str());
+    }
+  }
+
   // Print out the ODS information if requested.
   if (dumpODS)
     odsContext.print(llvm::errs());
@@ -68,11 +77,38 @@ processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
     pdlModule->print(os, OpPrintingFlags().enableDebugInfo());
     return success();
   }
-
   codegenPDLLToCPP(**module, *pdlModule, os);
   return success();
 }
 
+/// Create a dependency file for `-d` option.
+///
+/// This functionality is generally only for the benefit of the build system,
+/// and is modeled after the same option in TableGen.
+static LogicalResult
+createDependencyFile(StringRef outputFilename, StringRef dependencyFile,
+                     std::set<std::string> &includedFiles) {
+  if (outputFilename == "-") {
+    llvm::errs() << "error: the option -d must be used together with -o\n";
+    return failure();
+  }
+
+  std::string errorMessage;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      openOutputFile(dependencyFile, &errorMessage);
+  if (!outputFile) {
+    llvm::errs() << errorMessage << "\n";
+    return failure();
+  }
+
+  outputFile->os() << outputFilename << ":";
+  for (const auto &includeFile : includedFiles)
+    outputFile->os() << ' ' << includeFile;
+  outputFile->os() << "\n";
+  outputFile->keep();
+  return success();
+}
+
 int main(int argc, char **argv) {
   // FIXME: This is necessary because we link in TableGen, which defines its
   // options as static variables.. some of which overlap with our options.
@@ -110,6 +146,12 @@ int main(int argc, char **argv) {
                        clEnumValN(OutputType::CPP, "cpp",
                                   "generate a C++ source file containing the "
                                   "patterns for the input file")));
+  llvm::cl::opt<std::string> dependencyFilename(
+      "d", llvm::cl::desc("Dependency filename"),
+      llvm::cl::value_desc("filename"), llvm::cl::init(""));
+  llvm::cl::opt<bool> writeIfChanged(
+      "write-if-changed",
+      llvm::cl::desc("Only write to the output file if it changed"));
 
   llvm::InitLLVM y(argc, argv);
   llvm::cl::ParseCommandLineOptions(argc, argv, "PDLL Frontend");
@@ -123,28 +165,61 @@ int main(int argc, char **argv) {
     return 1;
   }
 
-  // Set up the output file.
-  std::unique_ptr<llvm::ToolOutputFile> outputFile =
-      openOutputFile(outputFilename, &errorMessage);
-  if (!outputFile) {
-    llvm::errs() << errorMessage << "\n";
-    return 1;
-  }
+  // If we are creating a dependency file, we'll also need to track what files
+  // get included during processing.
+  std::set<std::string> includedFilesStorage;
+  std::set<std::string> *includedFiles = nullptr;
+  if (!dependencyFilename.empty())
+    includedFiles = &includedFilesStorage;
 
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
+  std::string outputStr;
+  llvm::raw_string_ostream outputStrOS(outputStr);
   auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
                        raw_ostream &os) {
     return processBuffer(os, std::move(chunkBuffer), outputType, includeDirs,
-                         dumpODS);
+                         dumpODS, includedFiles);
   };
   if (splitInputFile) {
     if (failed(splitAndProcessBuffer(std::move(inputFile), processFn,
-                                     outputFile->os())))
+                                     outputStrOS)))
       return 1;
-  } else if (failed(processFn(std::move(inputFile), outputFile->os()))) {
+  } else if (failed(processFn(std::move(inputFile), outputStrOS))) {
     return 1;
   }
-  outputFile->keep();
+
+  // Write the output.
+  bool shouldWriteOutput = true;
+  if (writeIfChanged) {
+    // Only update the real output file if there are any differences. This
+    // prevents recompilation of all the files depending on it if there aren't
+    // any.
+    if (auto existingOrErr =
+            llvm::MemoryBuffer::getFile(outputFilename, /*IsText=*/true))
+      if (std::move(existingOrErr.get())->getBuffer() == outputStrOS.str())
+        shouldWriteOutput = false;
+  }
+
+  // Populate the output file if necessary.
+  if (shouldWriteOutput) {
+    std::unique_ptr<llvm::ToolOutputFile> outputFile =
+        openOutputFile(outputFilename, &errorMessage);
+    if (!outputFile) {
+      llvm::errs() << errorMessage << "\n";
+      return 1;
+    }
+    outputFile->os() << outputStrOS.str();
+    outputFile->keep();
+  }
+
+  // Always write the depfile, even if the main output hasn't changed. If it's
+  // missing, Ninja considers the output dirty.
+  if (!dependencyFilename.empty()) {
+    if (failed(createDependencyFile(outputFilename, dependencyFilename,
+                                    includedFilesStorage)))
+      return 1;
+  }
+
   return 0;
 }