#sdy add bytecode serialization utility

PiperOrigin-RevId: 730837065

shardy/tools/BUILD

@@ -15,3 +15,24 @@ cc_binary(
         "@llvm-project//mlir:QuantOps",
     ],
 )
+
+cc_binary(
+    name = "sdy_translate",
+    srcs = ["sdy_translate_main.cc"],
+    deps = [
+        "//shardy/dialect/sdy/ir:dialect",
+        "//shardy/dialect/sdy/ir:register",
+        "//shardy/dialect/sdy/transforms:passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:TranslateLib",
+    ],
+)

shardy/tools/sdy_translate_main.cc

@@ -0,0 +1,81 @@
+/* Copyright 2025 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// MLIR `translate` tool for allowing SDY dialect bytecode emission.
+//
+// Usage:
+//   sdy_translate <file.mlir> -serialize
+//   sdy_translate <file.mlir.bc> -deserialize
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
+#include "mlir/Dialect/Quant/IR/Quant.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"
+#include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
+#include "mlir/Tools/mlir-translate/Translation.h"
+#include "mlir/Transforms/Passes.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+#include "shardy/dialect/sdy/ir/register.h"
+
+namespace mlir {
+
+namespace {
+llvm::cl::opt<bool> stripDebuginfoOption(
+    "strip-debuginfo", llvm::cl::desc("Strip debug info from all operations"),
+    llvm::cl::init(false));
+
+void registerDialectsForSdy(DialectRegistry &registry) {
+  mlir::sdy::registerAllDialects(registry);
+  registry.insert<mlir::quant::QuantDialect>();
+}
+}  // namespace
+
+TranslateFromMLIRRegistration serializeRegistration(
+    "serialize", "Serialize SDY program into a portable artifact",
+    [](mlir::ModuleOp module, llvm::raw_ostream &os) -> llvm::LogicalResult {
+      if (stripDebuginfoOption) {
+        PassManager pm(module->getContext());
+        pm.addPass(createStripDebugInfoPass());
+        if (failed(pm.run(module)))
+          return module.emitError("failed to strip debuginfo");
+      }
+      const auto *producer = "SDY";
+      BytecodeWriterConfig writerConfig(producer);
+      return writeBytecodeToFile(module, os, writerConfig);
+    },
+    [](DialectRegistry &registry) { registerDialectsForSdy(registry); });
+
+TranslateToMLIRRegistration deserializeRegistration(
+    "deserialize", "Deserialize a portable artifact into a SDY program",
+    [](llvm::StringRef input, mlir::MLIRContext *context) {
+      context->loadDialect<sdy::SdyDialect>();
+      auto module = parseSourceString<ModuleOp>(input, context);
+      return module;
+    },
+    [](DialectRegistry &registry) { registerDialectsForSdy(registry); });
+
+}  // namespace mlir
+
+int main(int argc, char **argv) {
+  return mlir::asMainReturnCode(
+      mlir::mlirTranslateMain(argc, argv, "SDY transformation driver\n"));
+}