Add Check dialect to the MLIR adapter.

PiperOrigin-RevId: 647078265
google-ai-edge · Jun 26, 2024 · 9743466 · 9743466
1 parent 9e9d170
commit 9743466
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/builtin-adapter/BUILD b/src/builtin-adapter/BUILD
@@ -70,6 +70,7 @@ cc_library(
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:ShapeDialect",
         "@llvm-project//mlir:Support",
+        "@stablehlo//:check_ops",
         "@stablehlo//:chlo_ops",
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",

diff --git a/src/builtin-adapter/model_json_graph_convert.cc b/src/builtin-adapter/model_json_graph_convert.cc
@@ -52,6 +52,7 @@ limitations under the License.
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "stablehlo/dialect/VhloOps.h"
+#include "stablehlo/tests/CheckOps.h"
 #include "stablehlo/transforms/Passes.h"
 #include "tensorflow/cc/saved_model/reader.h"
 #include "status_macros.h"
@@ -352,7 +353,8 @@ absl::StatusOr<std::string> ConvertMlirToJson(const VisualizeConfig& config,
                   mlir::stablehlo::StablehloDialect, mlir::chlo::ChloDialect,
                   mlir::mhlo::MhloDialect, mlir::vhlo::VhloDialect,
                   mlir::func::FuncDialect, mlir::arith::ArithDialect,
-                  mlir::shape::ShapeDialect, mlir::scf::SCFDialect>();
+                  mlir::shape::ShapeDialect, mlir::scf::SCFDialect,
+                  mlir::stablehlo::check::CheckDialect>();
   mlir::MLIRContext context(registry);
   mlir::ParserConfig parser_config(&context);
   std::string model_content;