commit dc61fcc9d6fc9c237a4111d67a33fcc66b9a4ac2
author Scott Todd <scott.todd0@gmail.com> Mon May 20 13:40:18 2024 -0700
committer GitHub <noreply@github.com> Mon May 20 13:40:18 2024 -0700
tree dad17ece84c176b7637a170afa4453cd2d0f93ae
parent 2a2a4d020f7dc2e12f7768c322975068f41bbba1

Register ShapeDialect in StableHLO plugin. (#17444)

Fixes https://github.com/iree-org/iree/issues/16345

This fixes ``LLVM ERROR: Building op `shape.shape_of` but it isn't known
in this MLIRContext: the dialect may not be loaded or this operation
hasn't been added by the dialect`` errors. Something about how the pass
pipelines are constructed (dynamically?) made it so local registration
wasn't enough.

A few other dialects (tensor, arith, scf) appear to also be missing.
Might need to add those too.

compiler/plugins/input/StableHLO/BUILD.bazel

diff --git a/compiler/plugins/input/StableHLO/BUILD.bazel b/compiler/plugins/input/StableHLO/BUILD.bazel
index dfdb017..da1abae 100644
--- a/compiler/plugins/input/StableHLO/BUILD.bazel
+++ b/compiler/plugins/input/StableHLO/BUILD.bazel
@@ -28,6 +28,7 @@
         "@llvm-project//mlir:ConversionPasses",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
         "@llvm-project//mlir:Transforms",
         "@stablehlo//:chlo_ops",
         "@stablehlo//:stablehlo_ops",

compiler/plugins/input/StableHLO/CMakeLists.txt

diff --git a/compiler/plugins/input/StableHLO/CMakeLists.txt b/compiler/plugins/input/StableHLO/CMakeLists.txt
index 5f22917..1875e08 100644
--- a/compiler/plugins/input/StableHLO/CMakeLists.txt
+++ b/compiler/plugins/input/StableHLO/CMakeLists.txt
@@ -30,6 +30,7 @@
     ChloOps
     MLIRIR
     MLIRPass
+    MLIRShapeDialect
     MLIRTransforms
     StablehloOps
     VhloOps

compiler/plugins/input/StableHLO/PluginRegistration.cpp

diff --git a/compiler/plugins/input/StableHLO/PluginRegistration.cpp b/compiler/plugins/input/StableHLO/PluginRegistration.cpp
index 4f66ca21..df9f1a4 100644
--- a/compiler/plugins/input/StableHLO/PluginRegistration.cpp
+++ b/compiler/plugins/input/StableHLO/PluginRegistration.cpp
@@ -7,6 +7,7 @@
 #include "compiler/plugins/input/StableHLO/Conversion/Passes.h"
 #include "iree/compiler/PluginAPI/Client.h"
 #include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "stablehlo/dialect/ChloOps.h"
@@ -63,6 +64,7 @@
   }
 
   void onRegisterDialects(DialectRegistry &registry) override {
+    registry.insert<mlir::shape::ShapeDialect>();
     registry.insert<mlir::chlo::ChloDialect>();
     registry.insert<mlir::stablehlo::StablehloDialect>();
     registry.insert<mlir::vhlo::VhloDialect>();