Move softmax matcher and ErrorCheckingTrackingListener out of iree-dialects (#24466)

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -27,10 +27,6 @@ def __init__(self, repo_map: Dict[str, str]):
                 f"{iree_core_repo}//compiler/src/iree/compiler/API:CAPI": [
                     "IREECompilerCAPILib"
                 ],
-                # IREE llvm-external-projects
-                f"{iree_core_repo}//llvm-external-projects/iree-dialects:CAPI": [
-                    "IREEDialectsCAPI"
-                ],
                 # Disable all hard-coded codegen targets (they are expanded dynamically
                 # in CMake).
                 "@llvm-project//llvm:AArch64AsmParser": ["IREELLVMCPUTargetDeps"],

compiler/plugins/target/LLVMCPU/BUILD.bazel

@@ -47,7 +47,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/PluginAPI",
         "//compiler/src/iree/compiler/Utils",
-        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//llvm:AArch64AsmParser",
         "@llvm-project//llvm:AArch64CodeGen",
         "@llvm-project//llvm:ARMAsmParser",

compiler/plugins/target/LLVMCPU/CMakeLists.txt

@@ -31,7 +31,6 @@ iree_cc_library(
     ::LinkerTool
     ::StaticLibraryGenerator
     IREELLVMCPUTargetDeps
-    IREELinalgTransformDialect
     LLVMAnalysis
     LLVMBitReader
     LLVMBitWriter

compiler/src/iree/compiler/API/BUILD.bazel

@@ -37,7 +37,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/API/Internal:IREEOptToolEntryPoint",
         "//compiler/src/iree/compiler/API/Internal:IREEReduceToolEntryPoint",
         "//compiler/src/iree/compiler/API/Internal:LLDToolEntryPoint",
-        "//llvm-external-projects/iree-dialects:CAPI",
         "@llvm-project//mlir:CAPIAMDGPU",
         "@llvm-project//mlir:CAPIDebug",
         "@llvm-project//mlir:CAPIGPU",

compiler/src/iree/compiler/API/CMakeLists.txt

@@ -14,7 +14,6 @@ iree_cc_library(
   NAME
     StaticImpl
   DEPS
-    IREEDialectsCAPI
     MLIRCAPIAMDGPU
     MLIRCAPIDebug
     MLIRCAPIExportSMTLIB
@@ -82,7 +81,6 @@ set(_EXPORT_OBJECT_LIBS
   iree_compiler_API_Internal_IREEOptToolEntryPoint.objects
   iree_compiler_API_Internal_IREEReduceToolEntryPoint.objects
   iree_compiler_API_Internal_LLDToolEntryPoint.objects
-  obj.IREEDialectsCAPI
   obj.MLIRCAPIAMDGPU
   obj.MLIRCAPIDebug
   obj.MLIRCAPIExportSMTLIB

compiler/src/iree/compiler/API/api_exports.c

@@ -154,7 +154,6 @@ extern void ireeLinkRunMain();
 extern void ireeMlirLspServerRunMain();
 extern void ireeOptRunMain();
 extern void ireeReduceRunMain();
-extern void ireeRegisterTransformExtensions();
 extern void mlirAffineAddExprGet();
 extern void mlirAffineBinaryOpExprGetLHS();
 extern void mlirAffineBinaryOpExprGetRHS();
@@ -1330,7 +1329,6 @@ uintptr_t __iree_compiler_hidden_force_extern() {
   x += (uintptr_t)&ireeMlirLspServerRunMain;
   x += (uintptr_t)&ireeOptRunMain;
   x += (uintptr_t)&ireeReduceRunMain;
-  x += (uintptr_t)&ireeRegisterTransformExtensions;
   x += (uintptr_t)&mlirAffineAddExprGet;
   x += (uintptr_t)&mlirAffineBinaryOpExprGetLHS;
   x += (uintptr_t)&mlirAffineBinaryOpExprGetRHS;

compiler/src/iree/compiler/API/api_exports.def

@@ -144,7 +144,6 @@ EXPORTS
   ireeMlirLspServerRunMain
   ireeOptRunMain
   ireeReduceRunMain
-  ireeRegisterTransformExtensions
   mlirAffineAddExprGet
   mlirAffineBinaryOpExprGetLHS
   mlirAffineBinaryOpExprGetRHS

compiler/src/iree/compiler/API/api_exports.ld

@@ -145,7 +145,6 @@ VER_0 {
     ireeMlirLspServerRunMain;
     ireeOptRunMain;
     ireeReduceRunMain;
-    ireeRegisterTransformExtensions;
     mlirAffineAddExprGet;
     mlirAffineBinaryOpExprGetLHS;
     mlirAffineBinaryOpExprGetRHS;

compiler/src/iree/compiler/API/api_exports.macos.lst

@@ -143,7 +143,6 @@ _ireeLinkRunMain
 _ireeMlirLspServerRunMain
 _ireeOptRunMain
 _ireeReduceRunMain
-_ireeRegisterTransformExtensions
 _mlirAffineAddExprGet
 _mlirAffineBinaryOpExprGetLHS
 _mlirAffineBinaryOpExprGetRHS

compiler/src/iree/compiler/API/generate_exports.py

@@ -67,10 +67,6 @@
     "Dialect/PDL.h",
 ]
 
-IREE_DIALECTS_HEADER_FILES = [
-    "Dialects.h",
-]
-
 IREE_COMPILER_DIALECTS_HEADER_FILES = [
     "iree_codegen.h",
     "iree_gpu.h",
@@ -101,16 +97,6 @@ def main(repo_root: Path, api_root: Path):
     for local_name in LOCAL_HEADER_FILES:
         export_symbols.extend(collect_header_exports(api_root / local_name))
 
-    # Collect symbols from iree-dialects header files.
-    for local_name in IREE_DIALECTS_HEADER_FILES:
-        export_symbols.extend(
-            collect_header_exports(
-                repo_root
-                / "llvm-external-projects/iree-dialects/include/iree-dialects-c"
-                / local_name
-            )
-        )
-
     # Collect symbols from iree compiler dialect header files.
     for local_name in IREE_COMPILER_DIALECTS_HEADER_FILES:
         export_symbols.extend(

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

@@ -103,6 +103,7 @@ iree_compiler_cc_library(
         "EncodingUtils.cpp",
         "EraseDeadAllocAndStores.cpp",
         "EraseHALDescriptorTypeFromMemRef.cpp",
+        "ErrorCheckingTrackingListener.cpp",
         "FastMathPatterns.cpp",
         "FissionTransferOpsInControlFlow.cpp",
         "FlattenMemRefSubspan.cpp",
@@ -190,6 +191,7 @@ iree_compiler_cc_library(
         "CombineLayoutTransformation.h",
         "EmulateNarrowType.h",
         "EncodingUtils.h",
+        "ErrorCheckingTrackingListener.h",
         "FastMathPatterns.h",
         "PassUtils.h",
         "Passes.h",
@@ -234,7 +236,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/Dialect/Util/Transforms",
         "//compiler/src/iree/compiler/Utils",
-        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//llvm:Core",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AMDGPUDialect",
@@ -332,7 +333,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/TransformExtensions:LinalgExtExtensions",
         "//compiler/src/iree/compiler/Dialect/TensorExt/IR",
-        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:AffineUtils",
@@ -377,7 +377,6 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:DialectUtils",
         # TransformExtensions (needed for registration in the pass)
-        "//llvm-external-projects/iree-dialects:IREEDialectsTransforms",
         "//compiler/src/iree/compiler/Codegen/Common/TransformExtensions:CommonExtensions",
         "//compiler/src/iree/compiler/Codegen/LLVMCPU/TransformExtensions:LLVMCPUExtensions",
         "//compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions:LLVMGPUExtensions",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

@@ -56,6 +56,7 @@ iree_cc_library(
     "CombineLayoutTransformation.h"
     "EmulateNarrowType.h"
     "EncodingUtils.h"
+    "ErrorCheckingTrackingListener.h"
     "FastMathPatterns.h"
     "PassUtils.h"
     "Passes.h"
@@ -96,6 +97,7 @@ iree_cc_library(
     "EncodingUtils.cpp"
     "EraseDeadAllocAndStores.cpp"
     "EraseHALDescriptorTypeFromMemRef.cpp"
+    "ErrorCheckingTrackingListener.cpp"
     "FastMathPatterns.cpp"
     "FissionTransferOpsInControlFlow.cpp"
     "FlattenMemRefSubspan.cpp"
@@ -180,7 +182,6 @@ iree_cc_library(
   DEPS
     ::PassHeaders
     ::PassesIncGen
-    IREELinalgTransformDialect
     LLVMCore
     LLVMSupport
     MLIRAMDGPUDialect
@@ -287,8 +288,6 @@ iree_cc_library(
     ::Common
     ::PassHeaders
     ::PassesIncGen
-    IREEDialectsTransforms
-    IREELinalgTransformDialect
     LLVMSupport
     MLIRAMDGPUTransforms
     MLIRAffineDialect

compiler/src/iree/compiler/Codegen/Common/CommonDialectRegistration.cpp

@@ -4,7 +4,6 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.h"
 #include "iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
@@ -97,8 +96,7 @@ void registerTransformDialectTranslationDependentDialects(
   vector::registerBufferizableOpInterfaceExternalModels(registry);
 
   registry.addExtensions<
-      mlir::iree_compiler::IREE::LinalgExt::LinalgExtTransformOpsExtension,
-      transform_ext::StructuredTransformOpsExtension>();
+      mlir::iree_compiler::IREE::LinalgExt::LinalgExtTransformOpsExtension>();
   iree_compiler::registerTransformDialectCommonExtension(registry);
   iree_compiler::registerTransformDialectLLVMCPUExtension(registry);
   iree_compiler::registerTransformDialectLLVMGPUExtension(registry);

compiler/src/iree/compiler/Codegen/Common/ErrorCheckingTrackingListener.cpp

@@ -0,0 +1,40 @@
+// Copyright 2022 The IREE Authors
+//
+// Licensed 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 "iree/compiler/Codegen/Common/ErrorCheckingTrackingListener.h"
+
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+#define DEBUG_TYPE "iree-codegen-error-checking-tracking-listener"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+
+namespace mlir::iree_compiler {
+
+void ErrorCheckingTrackingListener::notifyPayloadReplacementNotFound(
+    Operation *op, ValueRange values, DiagnosedSilenceableFailure &&diag) {
+  // Certain ops can dropped safely.
+  if (isa<scf::ForOp>(op)) {
+    LLVM_DEBUG(DBGS() << "Silently dropping scf.for op mapping\n");
+    return;
+  }
+
+  SmallVector<Diagnostic> diags;
+  diag.takeDiagnostics(diags);
+  if (!status.succeeded()) {
+    status.takeDiagnostics(diags);
+  }
+  status = DiagnosedSilenceableFailure::silenceableFailure(std::move(diags));
+
+  status = emitSilenceableFailure(
+      getTransformOp(), "!!! tracking listener failed to find replacement op");
+  status.attachNote(op->getLoc()) << "replaced op";
+  for (Value v : values) {
+    status.attachNote(v.getLoc()) << "replacement value";
+  }
+}
+
+} // namespace mlir::iree_compiler

compiler/src/iree/compiler/Codegen/Common/ErrorCheckingTrackingListener.h

@@ -0,0 +1,48 @@
+// Copyright 2022 The IREE Authors
+//
+// Licensed 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 IREE_COMPILER_CODEGEN_COMMON_ERRORCHECKINGTRACKINGLISTENER_H_
+#define IREE_COMPILER_CODEGEN_COMMON_ERRORCHECKINGTRACKINGLISTENER_H_
+
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+
+namespace mlir::iree_compiler {
+
+/// A tracking listener for tensor IR that checks for payload replacement
+/// errors.
+class ErrorCheckingTrackingListener : public transform::TrackingListener {
+public:
+  using transform::TrackingListener::TrackingListener;
+
+  ~ErrorCheckingTrackingListener() override {
+    assert(status.succeeded() && "must check listener error state");
+  }
+
+  /// Return "true" if this tracking listener had a failure.
+  bool failed() const { return !status.succeeded(); }
+
+  /// Check and return the current error state of this listener. In case of a
+  /// failure state, only the most recent error is returned. Afterwards, resets
+  /// the error state.
+  DiagnosedSilenceableFailure checkAndResetError() {
+    DiagnosedSilenceableFailure result(std::move(status));
+    status = DiagnosedSilenceableFailure::success();
+    return result;
+  }
+
+private:
+  void
+  notifyPayloadReplacementNotFound(Operation *op, ValueRange values,
+                                   DiagnosedSilenceableFailure &&diag) override;
+
+  /// The error state of this listener. "Success" indicates that no error
+  /// happened so far. Otherwise, the status contains the most recent error.
+  DiagnosedSilenceableFailure status = DiagnosedSilenceableFailure::success();
+};
+
+} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_CODEGEN_COMMON_ERRORCHECKINGTRACKINGLISTENER_H_

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/BUILD.bazel

@@ -75,8 +75,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
-        "//llvm-external-projects/iree-dialects:IREEDialectsTransforms",
-        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:AffineUtils",

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CMakeLists.txt

@@ -31,8 +31,6 @@ iree_cc_library(
     "CommonExtensionsOps.cpp.inc"
   DEPS
     ::CommonExtensionsOpGen
-    IREEDialectsTransforms
-    IREELinalgTransformDialect
     LLVMSupport
     MLIRAffineDialect
     MLIRAffineUtils

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp

@@ -6,8 +6,7 @@
 
 #include "CommonExtensions.h"
 
-#include "iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.h"
-#include "iree-dialects/Transforms/TransformMatchers.h"
+#include "iree/compiler/Codegen/Common/ErrorCheckingTrackingListener.h"
 #include "iree/compiler/Codegen/Common/GPU/GPUPatterns.h"
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
@@ -48,6 +47,7 @@
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"

compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel

@@ -61,7 +61,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Dialect/PCF/ExternalInterfaces:ExternalModels",
         "//compiler/src/iree/compiler/Codegen/ExternalInterfaces:ExternalModels",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/TransformExtensions:LinalgExtExtensions",
-        "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgTransforms",
         "@llvm-project//mlir:SideEffectInterfaces",

compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt

@@ -25,7 +25,6 @@ iree_cc_library(
     ::TensorMaskingOpInterface
     ::UKernelOpInterface
     ::VectorizableOpInterface
-    IREELinalgTransformDialect
     MLIRAMDGPUDialect
     MLIRAffineDialect
     MLIRAffineTransformOps