[mlir] Start rewrite tool

Initial commit of a tool to help in textual rewrites of .mlir files. This tool
builds of of AsmParserState and is rather simple. Took some inspiration from
when I used clang's AST rewrites where I'd often treat it as a "localizing"
regex applicator in fallback cases, and started with that as functionality.
There though, one does have access to the lower level info than here, but still
a step up over sed over entire file.

This aims to be helpful (e.g., rewrite syntax including best effort inside
comments) rather than bulletproof tool. It may even be better suited under
utils than tools. And most of the rewrites would be rather short lived and
might never make it upstream (while the helpers of those rewrites may for
future rewrites).

The layering at the moment is not ideal as it is reusing the RewriteBuffer
class from clang's rewrite engine. So only optionally enabling where clang is
also enable. There doesn't seem to be anything clang specific there (the dep
does pull in more dependencies than ideal, but leaving both refactorings).

Additionally started it as a single file to prototype more easily, planning to
refactor later to include and libs for out of file usage.

Author: jpienaar

mlir/CMakeLists.txt

@@ -285,3 +285,8 @@ endif()
 if(MLIR_STANDALONE_BUILD)
   llvm_distribution_add_targets()
 endif()
+
+# FIXME: Currently depends on utility functions inside clang.
+if ("clang" IN_LIST LLVM_ENABLE_PROJECTS)
+  set(MLIR_ENABLE_REWRITE ON CACHE BOOL "mlir-rewrite enabled")
+endif()

mlir/docs/Tools/mlir-rewrite.md

@@ -0,0 +1,29 @@
+# mlir-rewrite
+
+Tool to simplify rewriting .mlir files. There are a couple of build in rewrites
+discussed below along with usage.
+
+Note: This is still in very early stage. Its so early its less a tool than a
+growing collection of useful functions: to use its best to do what's needed on
+a brance by just hacking it (dialects registered, rewrites etc) to say help
+ease a rename, upstream useful utility functions, point to ease others
+migrating, and then bin eventually. Once there are actually useful parts it
+should be refactored same as mlir-opt.
+
+[TOC]
+
+## simple-rename
+
+Rename per op given a substring to a target. The match and replace uses LLVM's
+regex sub for the match and replace while the op-name is matched via regular
+string comparison. E.g.,
+
+```
+mlir-rewrite input.mlir -o output.mlir --simple-rename \
+   --simple-rename-op-name="test.concat" --simple-rename-match="axis" \
+                                         --simple-rename-replace="bxis"
+```
+
+to replace `axis` substring in the text of the range corresponding to
+`test.concat` ops with `bxis`.
+

mlir/test/CMakeLists.txt

@@ -197,6 +197,13 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   )
 endif()
 
+# FIXME: Currently depends on utility functions inside clang.
+if ("clang" IN_LIST LLVM_ENABLE_PROJECTS)
+  list(APPEND MLIR_TEST_DEPENDS
+    mlir-rewrite
+  )
+endif()
+
 # This target can be used to just build the dependencies
 # for the check-mlir target without executing the tests.
 # This is useful for bots when splitting the build step

mlir/test/lit.cfg.py

@@ -144,6 +144,10 @@ def add_runtime(name):
         )
     )
 
+if config.enable_mlir_rewrite:
+    tools.extend(["mlir-rewrite"])
+    config.available_features.add('mlir-rewrite')
+
 # The following tools are optional
 tools.extend(
     [

mlir/test/lit.site.cfg.py.in

@@ -23,6 +23,7 @@ config.mlir_obj_root = "@MLIR_BINARY_DIR@"
 config.mlir_tools_dir = "@MLIR_TOOLS_DIR@"
 config.mlir_cmake_dir = "@MLIR_CMAKE_DIR@"
 config.mlir_lib_dir = "@MLIR_LIB_DIR@"
+config.enable_mlir_rewrite = "@MLIR_ENABLE_REWRITE@"
 
 config.build_examples = @LLVM_BUILD_EXAMPLES@
 config.run_cuda_tests = @MLIR_ENABLE_CUDA_CONVERSIONS@

mlir/test/mlir-rewrite/simple.mlir

@@ -0,0 +1,12 @@
+// RUN: mlir-opt %s | mlir-rewrite --simple-rename --simple-rename-op-name="test.concat" --simple-rename-match="axis" --simple-rename-replace="bxis" | FileCheck %s -check-prefix=RENAME
+// RUN: mlir-opt %s | mlir-rewrite --mark-ranges | FileCheck %s -check-prefix=RANGE
+// Note: running through mlir-opt to just strip out comments & avoid self matches.
+// REQUIRES: mlir-rewrite
+
+func.func @two_dynamic_one_direct_shape(%arg0: tensor<?x4x?xf32>, %arg1: tensor<2x4x?xf32>) -> tensor<?x4x?xf32> {
+  // RENAME: "test.concat"({{.*}}) {bxis = 0 : i64}
+  // RANGE: 《%{{.*}} = 〖"test.concat"〗({{.*}}) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>》
+  %5 = "test.concat"(%arg0, %arg1) {axis = 0 : i64} : (tensor<?x4x?xf32>, tensor<2x4x?xf32>) -> tensor<?x4x?xf32>
+  return %5 : tensor<?x4x?xf32>
+}
+

mlir/tools/CMakeLists.txt

@@ -15,3 +15,8 @@ add_subdirectory(tblgen-to-irdl)
 if(MLIR_ENABLE_EXECUTION_ENGINE)
   add_subdirectory(mlir-cpu-runner)
 endif()
+
+# FIXME: Currently depends on utility functions inside clang.
+if ("clang" IN_LIST LLVM_ENABLE_PROJECTS)
+  add_subdirectory(mlir-rewrite)
+endif()

mlir/tools/mlir-rewrite/CMakeLists.txt

@@ -0,0 +1,37 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+set(LLVM_LINK_COMPONENTS
+  Support
+  )
+
+set(LIBS
+  ${dialect_libs}
+  ${test_libs}
+
+  clangRewrite
+  MLIRAffineAnalysis
+  MLIRAnalysis
+  MLIRCastInterfaces
+  MLIRDialect
+  MLIROptLib
+  MLIRParser
+  MLIRPass
+  MLIRTransforms
+  MLIRTransformUtils
+  MLIRSupport
+  MLIRIR
+  )
+
+include_directories(../../../clang/include)
+
+add_mlir_tool(mlir-rewrite
+  mlir-rewrite.cpp
+
+  DEPENDS
+  ${LIBS}
+  SUPPORT_PLUGINS
+  )
+target_link_libraries(mlir-rewrite PRIVATE ${LIBS})
+llvm_update_compile_flags(mlir-rewrite)
+
+mlir_check_all_link_libraries(mlir-rewrite)
+export_executable_symbols_for_plugins(mlir-rewrite)