[Mlir-commits] [mlir] [mlir] introduce debug transform dialect extension (PR #77595)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Wed Jan 10 04:36:17 PST 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

<details>
<summary>Changes</summary>

Introduce a new extension for simple print-debugging of the transform dialect scripts. The initial version of this extension consists of two ops that are printing the payload objects associated with transform dialect values. Similar ops were already available in the test extenion and several downstream projects, and were extensively used for testing.

---

Patch is 128.84 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/77595.diff


60 Files Affected:

- (modified) mlir/docs/Dialects/Transform.md (+4) 
- (modified) mlir/docs/Tutorials/transform/Ch1.md (+7-7) 
- (modified) mlir/docs/Tutorials/transform/Ch4.md (+2-2) 
- (modified) mlir/include/mlir/Dialect/Transform/CMakeLists.txt (+1) 
- (added) mlir/include/mlir/Dialect/Transform/DebugExtension/CMakeLists.txt (+6) 
- (added) mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtension.h (+22) 
- (added) mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h (+23) 
- (added) mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.td (+67) 
- (modified) mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h (+9-8) 
- (modified) mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td (+5) 
- (modified) mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h (+5) 
- (modified) mlir/include/mlir/InitAllExtensions.h (+2) 
- (modified) mlir/lib/Dialect/Transform/CMakeLists.txt (+1) 
- (added) mlir/lib/Dialect/Transform/DebugExtension/CMakeLists.txt (+11) 
- (added) mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp (+35) 
- (added) mlir/lib/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp (+70) 
- (modified) mlir/test/Dialect/Linalg/match-ops-interpreter.mlir (+49-49) 
- (modified) mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir (+5-5) 
- (modified) mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir (+1-1) 
- (modified) mlir/test/Dialect/Linalg/transform-op-match.mlir (+11-11) 
- (modified) mlir/test/Dialect/Linalg/transform-op-multitile-sizes.mlir (+3-3) 
- (modified) mlir/test/Dialect/Linalg/transform-op-pad.mlir (+1-1) 
- (modified) mlir/test/Dialect/MemRef/extract-address-computations.mlir (+1-1) 
- (modified) mlir/test/Dialect/MemRef/transform-ops.mlir (+5-5) 
- (modified) mlir/test/Dialect/SCF/transform-ops.mlir (+5-5) 
- (modified) mlir/test/Dialect/SparseTensor/transform-ops.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/expensive-checks.mlir (+3-3) 
- (modified) mlir/test/Dialect/Transform/include/test-interpreter-external-concurrent-source.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/include/test-interpreter-external-source.mlir (+2-2) 
- (modified) mlir/test/Dialect/Transform/include/test-interpreter-external-symbol-def-invalid.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/include/test-interpreter-library-invalid/definitions-invalid.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/include/test-interpreter-library/definitions-self-contained.mlir (+8-8) 
- (modified) mlir/test/Dialect/Transform/infer-effects.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/multi-arg-top-level-ops.mlir (+4-4) 
- (modified) mlir/test/Dialect/Transform/multi-arg-top-level-params.mlir (+2-2) 
- (modified) mlir/test/Dialect/Transform/multi-arg-top-level-values.mlir (+2-2) 
- (modified) mlir/test/Dialect/Transform/ops-invalid.mlir (+3-3) 
- (modified) mlir/test/Dialect/Transform/test-interpreter-debug.mlir (+2-2) 
- (modified) mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl-invalid.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/test-interpreter-external-symbol-decl.mlir (+7-7) 
- (modified) mlir/test/Dialect/Transform/test-interpreter-multiple-top-level-ops.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/test-interpreter.mlir (+77-77) 
- (modified) mlir/test/Dialect/Transform/test-loop-transforms.mlir (+4-4) 
- (modified) mlir/test/Dialect/Transform/test-pattern-application.mlir (+4-4) 
- (modified) mlir/test/Dialect/Transform/test-pdl-extension.mlir (+1-1) 
- (modified) mlir/test/Dialect/Transform/test-repro-dump.mlir (+1-1) 
- (modified) mlir/test/Examples/transform/Ch1/invalidation-1.mlir (+2-2) 
- (modified) mlir/test/Examples/transform/Ch1/invalidation-2.mlir (+1-1) 
- (modified) mlir/test/Examples/transform/Ch4/features.mlir (+1-1) 
- (modified) mlir/test/Examples/transform/Ch4/multiple.mlir (+3-3) 
- (modified) mlir/test/Examples/transform/Ch4/sequence.mlir (+2-2) 
- (modified) mlir/test/Integration/Dialect/Transform/match_batch_matmul.mlir (+7-7) 
- (modified) mlir/test/Integration/Dialect/Transform/match_matmul.mlir (+6-6) 
- (modified) mlir/test/Integration/Dialect/Transform/match_reduction.mlir (+7-7) 
- (modified) mlir/test/lib/Dialect/Transform/CMakeLists.txt (+1) 
- (modified) mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp (-68) 
- (modified) mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td (-36) 
- (modified) mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp (+3-2) 
- (modified) mlir/unittests/Dialect/Transform/CMakeLists.txt (+1) 
- (modified) mlir/unittests/Dialect/Transform/Preload.cpp (+3-5) 


``````````diff
diff --git a/mlir/docs/Dialects/Transform.md b/mlir/docs/Dialects/Transform.md
index 67b4ee53c1b573..768de9561b9515 100644
--- a/mlir/docs/Dialects/Transform.md
+++ b/mlir/docs/Dialects/Transform.md
@@ -423,6 +423,10 @@ ops rather than having the methods directly act on the payload IR.
 
 [include "Dialects/BufferizationTransformOps.md"]
 
+## Debug Transform Operations
+
+[include "Dialects/DebugExtensionOps.md"]
+
 ## Func Transform Operations
 
 [include "Dialects/FuncTransformOps.md"]
diff --git a/mlir/docs/Tutorials/transform/Ch1.md b/mlir/docs/Tutorials/transform/Ch1.md
index 8d9dfac18b5315..7a299a48600b8f 100644
--- a/mlir/docs/Tutorials/transform/Ch1.md
+++ b/mlir/docs/Tutorials/transform/Ch1.md
@@ -69,9 +69,9 @@ transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op,
      %arg1: !transform.op<"linalg.matmul">,
      %arg2: !transform.op<"linalg.elemwise_binary">):
-  transform.test_print_remark_at_operand %arg1, "matmul"
+  transform.debug.emit_remark_at %arg1, "matmul"
       : !transform.op<"linalg.matmul">
-  transform.test_print_remark_at_operand %arg2, "elemwise_binaries"
+  transform.debug.emit_remark_at %arg2, "elemwise_binaries"
       : !transform.op<"linalg.elemwise_binary">
   transform.yield
 }
@@ -180,7 +180,7 @@ transform.sequence failures(propagate) {
       : (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
 
   // This is trying to use an invalidated handle leading to undefined behavior.
-  transform.test_print_remark_at_operand %arg1, "remark" : !transform.op<"linalg.matmul">
+  transform.debug.emit_remark_at %arg1, "remark" : !transform.op<"linalg.matmul">
   transform.yield
 }
 ```
@@ -197,7 +197,7 @@ $ mlir-opt matmul.mlir --pass-pipeline="
 
 ```sh
 matmul.mlir:28:3: error: op uses a handle invalidated by a previously executed transform op
-  transform.test_print_remark_at_operand %mm, "elemwise_binaries" : !transform.any_op
+  transform.debug.emit_remark_at %mm, "elemwise_binaries" : !transform.any_op
   ^
 matmul.mlir:26:9: note: handle to invalidated ops
   %mm = transform.cast %matmul : !transform.op<"linalg.matmul"> to !transform.any_op
@@ -224,7 +224,7 @@ transform.sequence failures(propagate) {
 
   // Consuming an operand invalidates the consumed handle and any other handle that is
   // associated with the same payload operations, or payload operations nested in them.
-  transform.test_print_remark_at_operand %casted, "remark"
+  transform.debug.emit_remark_at %casted, "remark"
     : !transform.any_op
   transform.yield
 }
@@ -234,7 +234,7 @@ Both `%arg1` and `%casted` reference the same payload operation. Extending the r
 
 ```sh
 matmul.mlir:28:3: error: op uses a handle invalidated by a previously executed transform op
-  transform.test_print_remark_at_operand %matmul, "elemwise_binaries" : !transform.op<"linalg.matmul">
+  transform.debug.emit_remark_at %matmul, "elemwise_binaries" : !transform.op<"linalg.matmul">
   ^
 matmul.mlir:21:29: note: handle to invalidated ops
 ^bb0(%root: !transform.any_op, %matmul: !transform.op<"linalg.matmul">, %elemwise: !transform.op<"linalg.elemwise_binary">):
@@ -358,7 +358,7 @@ Attempting to access the fusion result after outlining produces the following er
 
 ```sh
 test/Examples/transform/Ch1/invalidation-2.mlir:109:3: error: op uses a handle invalidated by a previously executed transform op
-  transform.test_print_remark_at_operand %outline_target, "outlined loop" : !transform.any_op
+  transform.debug.emit_remark_at %outline_target, "outlined loop" : !transform.any_op
   ^
 test/Examples/transform/Ch1/invalidation-2.mlir:102:25: note: handle to invalidated ops
   %outline_target, %_ = transform.structured.tile_using_forall %tiled_2 tile_sizes [1]
diff --git a/mlir/docs/Tutorials/transform/Ch4.md b/mlir/docs/Tutorials/transform/Ch4.md
index 77c36eab343d24..9c9aba19d5745c 100644
--- a/mlir/docs/Tutorials/transform/Ch4.md
+++ b/mlir/docs/Tutorials/transform/Ch4.md
@@ -110,13 +110,13 @@ module @transforms attributes { transform.with_named_sequence } {
   // This is a rewriter sequence.
   transform.named_sequence @print_elemwise(
       %elemwise_binary: !transform.any_op {transform.readonly}) {
-    transform.test_print_remark_at_operand
+    transform.debug.emit_remark_at
       %elemwise_binary, "elementwise binary" : !transform.any_op
     transform.yield
   }
   transform.named_sequence @print_matmul(
       %matmul: !transform.any_op {transform.readonly}) {
-    transform.test_print_remark_at_operand %matmul, "matmul" : !transform.any_op
+    transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op
     transform.yield
   }
 }
diff --git a/mlir/include/mlir/Dialect/Transform/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
index d6c5c975c2e93c..ef44ab891cc5c9 100644
--- a/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/Transform/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(DebugExtension)
 add_subdirectory(IR)
 add_subdirectory(LoopExtension)
 add_subdirectory(PDLExtension)
diff --git a/mlir/include/mlir/Dialect/Transform/DebugExtension/CMakeLists.txt b/mlir/include/mlir/Dialect/Transform/DebugExtension/CMakeLists.txt
new file mode 100644
index 00000000000000..2b49faef7f00b6
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/DebugExtension/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS DebugExtensionOps.td)
+mlir_tablegen(DebugExtensionOps.h.inc -gen-op-decls)
+mlir_tablegen(DebugExtensionOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRTransformDialectDebugExtensionOpsIncGen)
+
+add_mlir_doc(DebugExtensionOps DebugExtensionOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtension.h b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtension.h
new file mode 100644
index 00000000000000..5052f763335e83
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtension.h
@@ -0,0 +1,22 @@
+//===- DebugExtension.h - Debug extension for Transform dialect -*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSION_H
+#define MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSION_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace transform {
+/// Registers the debug extension of the Transform dialect in the given
+/// registry.
+void registerDebugExtension(DialectRegistry &dialectRegistry);
+} // namespace transform
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSION_H
diff --git a/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h
new file mode 100644
index 00000000000000..d660393a71e0a4
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h
@@ -0,0 +1,23 @@
+//===- DebugExtensionOps.h - Debug ext. for Transform dialect ---*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSIONOPS_H
+#define MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSIONOPS_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DEBUGEXTENSIONOPS_H
diff --git a/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.td b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.td
new file mode 100644
index 00000000000000..16e2a39044b496
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.td
@@ -0,0 +1,67 @@
+//===- DebugExtensionOps.td - Transform Debug extension ----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines operations of the transform dialect extension for debugging transform
+// scripts.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DBEUGEXTENSIONOPS
+#define MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DBEUGEXTENSIONOPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Transform/IR/MatchInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+
+def DebugEmitRemarkAtOp : TransformDialectOp<"debug.emit_remark_at",
+  [MatchOpInterface,
+   DeclareOpInterfaceMethods<TransformOpInterface>,
+   MemoryEffectsOpInterface, NavigationTransformOpTrait]> {
+  let summary = "Print a message as diagnostic remark attached to payload";
+  let description = [{
+    This operation emits a diagnostic remark with the given message at the
+    location of each payload object associated with the argument. The argument
+    may be an operation or a value handle.
+
+    This operation always succeeds.
+  }];
+
+  let arguments = (ins
+    Transform_AnyHandleType:$at,
+    StrAttr:$message);
+  let assemblyFormat = "$at `,` $message attr-dict `:` type($at)";
+}
+
+def DebugEmitParamAsRemarkOp 
+  : TransformDialectOp<"debug.emit_param_as_remark",
+    [MatchOpInterface,
+     DeclareOpInterfaceMethods<TransformOpInterface>,
+     MemoryEffectsOpInterface, NavigationTransformOpTrait]> {
+  let summary = "Prints the parameter as a diagnostic remark";
+  let description = [{
+    This operation emits a diagnostic remark containing the string form of the
+    attributes associated with the parameter provided as attribute. It takes
+    as optional arguments:
+      - an additional message text to prepend;
+      - a handle pointing to operations the location of which will be used to
+        emit the diagnostic; if multiple operations are associated, the
+        diagnostic is emitted for all of their respective locations.
+
+    This operation always succeeds.
+  }];
+
+  let arguments = (ins TransformParamTypeInterface:$param,
+                       Optional<TransformHandleTypeInterface>:$anchor,
+                       OptionalAttr<StrAttr>:$message);
+  let assemblyFormat = "$param (`,` $message^)?  (`at` $anchor^)?"
+                       "attr-dict `:` type($param) (`,` type($anchor)^)?";
+}
+
+#endif // MLIR_DIALECT_TRANSFORM_DEBUGEXTENSION_DBEUGEXTENSIONOPS
diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
index 23871cc16d87de..c2e3cde8ebc69f 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
@@ -1282,9 +1282,9 @@ class FunctionalStyleTransformOpTrait
   }
 };
 
-/// Trait implementing the MemoryEffectOpInterface for single-operand
-/// single-result operations that use their operand without consuming and
-/// without modifying the Payload IR to produce a new handle.
+/// Trait implementing the MemoryEffectOpInterface for operations that use their
+/// operands without consuming and without modifying the Payload IR to
+/// potentially produce new handles.
 template <typename OpTy>
 class NavigationTransformOpTrait
     : public OpTrait::TraitBase<OpTy, NavigationTransformOpTrait> {
@@ -1294,15 +1294,16 @@ class NavigationTransformOpTrait
   void getEffects(SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
     onlyReadsHandle(this->getOperation()->getOperands(), effects);
     producesHandle(this->getOperation()->getResults(), effects);
-    onlyReadsPayload(effects);
+    if (llvm::any_of(this->getOperation()->getOperandTypes(), [](Type t) {
+          return isa<TransformHandleTypeInterface,
+                     TransformValueHandleTypeInterface>(t);
+        })) {
+      onlyReadsPayload(effects);
+    }
   }
 
   /// Checks that the op matches the expectation of this trait.
   static LogicalResult verifyTrait(Operation *op) {
-    static_assert(OpTy::template hasTrait<OpTrait::OneOperand>(),
-                  "expected single-operand op");
-    static_assert(OpTy::template hasTrait<OpTrait::OneResult>(),
-                  "expected single-result op");
     if (!op->getName().getInterface<MemoryEffectOpInterface>()) {
       op->emitError() << "NavigationTransformOpTrait should only be attached "
                          "to ops that implement MemoryEffectOpInterface";
diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td
index df04c52bb9cb8e..f29efaee620d84 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td
@@ -182,6 +182,11 @@ def Transform_AnyHandleOrParamType
              TransformValueHandleTypeInterface.predicate]>,
          "any transform handle or parameter">;
 
+def Transform_AnyHandleType
+  : Type<Or<[TransformHandleTypeInterface.predicate,
+             TransformValueHandleTypeInterface.predicate]>,
+         "any transform handle">;
+
 def FunctionalStyleTransformOpTrait
     : NativeOpTrait<"FunctionalStyleTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
diff --git a/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h b/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h
index 7a8ed2075ef12e..b62e6193d97aff 100644
--- a/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h
+++ b/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtension.h
@@ -6,6 +6,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSION_H
+#define MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSION_H
+
 namespace mlir {
 class DialectRegistry;
 
@@ -14,3 +17,5 @@ namespace transform {
 void registerLoopExtension(DialectRegistry &dialectRegistry);
 } // namespace transform
 } // namespace mlir
+
+#endif // MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSION_H
diff --git a/mlir/include/mlir/InitAllExtensions.h b/mlir/include/mlir/InitAllExtensions.h
index c04ce850fb96f4..7708ca5571de3b 100644
--- a/mlir/include/mlir/InitAllExtensions.h
+++ b/mlir/include/mlir/InitAllExtensions.h
@@ -34,6 +34,7 @@
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtension.h"
 #include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
 #include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
@@ -75,6 +76,7 @@ inline void registerAllExtensions(DialectRegistry &registry) {
   scf::registerTransformDialectExtension(registry);
   sparse_tensor::registerTransformDialectExtension(registry);
   tensor::registerTransformDialectExtension(registry);
+  transform::registerDebugExtension(registry);
   transform::registerLoopExtension(registry);
   transform::registerPDLExtension(registry);
   vector::registerTransformDialectExtension(registry);
diff --git a/mlir/lib/Dialect/Transform/CMakeLists.txt b/mlir/lib/Dialect/Transform/CMakeLists.txt
index 6898d81df7ca63..ed05194b72990b 100644
--- a/mlir/lib/Dialect/Transform/CMakeLists.txt
+++ b/mlir/lib/Dialect/Transform/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(DebugExtension)
 add_subdirectory(IR)
 add_subdirectory(LoopExtension)
 add_subdirectory(PDLExtension)
diff --git a/mlir/lib/Dialect/Transform/DebugExtension/CMakeLists.txt b/mlir/lib/Dialect/Transform/DebugExtension/CMakeLists.txt
new file mode 100644
index 00000000000000..38aca4c7340e97
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/DebugExtension/CMakeLists.txt
@@ -0,0 +1,11 @@
+add_mlir_dialect_library(MLIRTransformDebugExtension
+  DebugExtension.cpp
+  DebugExtensionOps.cpp
+
+  DEPENDS
+  MLIRTransformDialectDebugExtensionOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRTransformDialect
+)
diff --git a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
new file mode 100644
index 00000000000000..2901a35596f0b4
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
@@ -0,0 +1,35 @@
+//===- DebugExtension.cpp - Debug extension for the Transform dialect
+//-------===//
+//
+// 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/Transform/DebugExtension/DebugExtension.h"
+
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/IR/DialectRegistry.h"
+
+using namespace mlir;
+
+namespace {
+/// Debug extension of the Transform dialect. This provides operations for
+/// debugging transform dialect scripts.
+class DebugExtension
+    : public transform::TransformDialectExtension<DebugExtension> {
+public:
+  void init() {
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+void mlir::transform::registerDebugExtension(DialectRegistry &dialectRegistry) {
+  dialectRegistry.addExtensions<DebugExtension>();
+}
diff --git a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp
new file mode 100644
index 00000000000000..4b3e28e4313c64
--- /dev/null
+++ b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp
@@ -0,0 +1,70 @@
+//===- DebugExtensionOps.cpp - Debug extension for the Transform dialect --===//
+//
+// 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/Transform/DebugExtension/DebugExtensionOps.h"
+
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/IR/OpImplementation.h"
+
+using namespace mlir;
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp.inc"
+
+DiagnosedSilenceableFailure
+transform::DebugEmitRemarkAtOp::apply(transform::TransformRewriter &rewriter,
+                                      transform::TransformResults &results,
+                                      transform::TransformState &state) {
+  if (getAt().getType().isa<TransformHandleTypeInterface>()) {
+    auto payload = state.getPayloadOps(getAt());
+    for (Operation *op : payload)
+      op->emitRemark() << getMessage();
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  assert(
+      getAt().getType().isa<transform::TransformValueHandleTypeInterface>() &&
+      "unhandled kind of transform type");
+
+  auto describeValue = [](Diagnostic &os, Value value) {
+    os << "value handle points to ";
+    if (auto arg = llvm::dyn_cast<BlockArgument>(value)) {
+      os << "a block argument #" << arg.getArgNumber(...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/77595


More information about the Mlir-commits mailing list