[mlir] Start rewrite tool (#77668)

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/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

@@ -115,6 +115,7 @@ set(MLIR_TEST_DEPENDS
   mlir-opt
   mlir-query
   mlir-reduce
+  mlir-rewrite
   mlir-tblgen
   mlir-translate
   tblgen-lsp-server

mlir/test/mlir-rewrite/simple.mlir

@@ -0,0 +1,11 @@
+// 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.
+
+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

@@ -4,6 +4,7 @@ add_subdirectory(mlir-parser-fuzzer)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-query)
 add_subdirectory(mlir-reduce)
+add_subdirectory(mlir-rewrite)
 add_subdirectory(mlir-shlib)
 add_subdirectory(mlir-spirv-cpu-runner)
 add_subdirectory(mlir-translate)

mlir/tools/mlir-rewrite/CMakeLists.txt

@@ -0,0 +1,35 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+set(LLVM_LINK_COMPONENTS
+  Support
+  )
+
+set(LIBS
+  ${dialect_libs}
+  ${test_libs}
+
+  MLIRAffineAnalysis
+  MLIRAnalysis
+  MLIRCastInterfaces
+  MLIRDialect
+  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)