#sdy Duplicate addSdyRoundTripImportPipeline and

registerSdyRoundTripImportPipeline to remove xla dependence. This is a temporary solution before we have stablehlo serialization and deserialization support. PiperOrigin-RevId: 730920864
openxla · Feb 25, 2025 · 82f3d21 · 82f3d21
1 parent 1be3888
commit 82f3d21
Show file tree

Hide file tree

Showing 26 changed files with 2,763 additions and 0 deletions.
diff --git a/shardy/round_trip_import/BUILD b/shardy/round_trip_import/BUILD
@@ -0,0 +1,189 @@
+package(
+    default_applicable_licenses = ["//third_party/openxla:package_license"],
+    default_restricted_to = [
+        "//buildenv/target:non_prod",
+        "//buildenv/target:prod",
+    ],
+    default_visibility = ["//shardy:friends"],
+)
+
+cc_library(
+    name = "constants",
+    hdrs = ["constants.h"],
+    deps = ["@llvm-project//llvm:Support"],
+)
+
+cc_library(
+    name = "utils",
+    srcs = ["utils.cc"],
+    hdrs = ["utils.h"],
+    deps = [
+        ":constants",
+        "//shardy/dialect/sdy/ir:dialect",
+        "//shardy/dialect/sdy/ir:register",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AsmParser",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "import_sdy_custom_calls",
+    srcs = ["import_sdy_custom_calls.cc"],
+    hdrs = ["import_sdy_custom_calls.h"],
+    deps = [
+        ":constants",
+        ":utils",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "import_backend_func_calls",
+    srcs = ["import_backend_func_calls.cc"],
+    hdrs = ["import_backend_func_calls.h"],
+    deps = [
+        ":constants",
+        ":utils",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "import_constants",
+    srcs = ["import_constants.cc"],
+    hdrs = ["import_constants.h"],
+    deps = [
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "open_while_free_vars_sharding",
+    srcs = ["open_while_free_vars_sharding.cc"],
+    hdrs = ["open_while_free_vars_sharding.h"],
+    deps = [
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "import_shardy_attrs",
+    srcs = ["import_shardy_attrs.cc"],
+    hdrs = ["import_shardy_attrs.h"],
+    deps = [
+        "constants",
+        ":utils",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AsmParser",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "remove_size_one_axes",
+    srcs = ["remove_size_one_axes.cc"],
+    hdrs = ["remove_size_one_axes.h"],
+    deps = [
+        "//shardy/dialect/sdy/ir:dialect",
+        "//shardy/dialect/sdy/transforms/common:sharding_walker",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "import_callback_custom_calls",
+    srcs = ["import_callback_custom_calls.cc"],
+    hdrs = ["import_callback_custom_calls.h"],
+    deps = [
+        ":utils",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "shard_map_import",
+    srcs = ["shard_map_import.cc"],
+    hdrs = ["shard_map_import.h"],
+    deps = [
+        ":constants",
+        ":utils",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "pipelines",
+    srcs = ["pipelines.cc"],
+    hdrs = ["pipelines.h"],
+    deps = [
+        ":import_backend_func_calls",
+        ":import_callback_custom_calls",
+        ":import_constants",
+        ":import_sdy_custom_calls",
+        ":import_shardy_attrs",
+        ":open_while_free_vars_sharding",
+        ":remove_size_one_axes",
+        ":shard_map_import",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:stablehlo_passes",
+    ],
+)
diff --git a/shardy/round_trip_import/README.md b/shardy/round_trip_import/README.md
@@ -0,0 +1,8 @@
+#Shardy round_trip_import
+
+This directory is only temporarily exists and will be removed once stablehlo
+supports deserialization/serialization with different dialects.
+
+This directory duplicate the xla/service/spmd/shardy implementation for
+addSdyRoundTripImportPipeline and registerSdyRoundTripImportPipeline, to remove
+the xla dependency.
diff --git a/shardy/round_trip_import/constants.h b/shardy/round_trip_import/constants.h
@@ -0,0 +1,141 @@
+/* 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.
+==============================================================================*/
+
+#ifndef SHARDY_ROUND_TRIP_IMPORT_CONSTANTS_H_
+#define SHARDY_ROUND_TRIP_IMPORT_CONSTANTS_H_
+
+#include "llvm/ADT/StringRef.h"
+
+namespace xla {
+namespace sdy {
+
+// TODO(b/393406516): [Shardy] Change Attributes names from *mhlo* to
+// *stablehlo*
+inline constexpr llvm::StringRef kMhloAttributesAttr = "mhlo.attributes";
+
+// The attribute name for xla::HloSharding.
+inline constexpr llvm::StringRef kXlaShardingAttr = "mhlo.sharding";
+
+// The target name of the Sharding custom call.
+inline constexpr llvm::StringRef kShardingCustomCallTargetName = "Sharding";
+
+// The target name of the SPMDFullToShardShape custom call.
+inline constexpr llvm::StringRef kSPMDFullToShardShapeCallTargetName =
+    "SPMDFullToShardShape";
+
+// The target name of the SPMDShardToFullShape custom call.
+inline constexpr llvm::StringRef kSPMDShardToFullShapeCallTargetName =
+    "SPMDShardToFullShape";
+
+// The target name of the Python CPU callback custom call.
+inline constexpr llvm::StringRef kPythonCpuCallbackCustomCallTargetName =
+    "xla_python_cpu_callback";
+
+// The target name of the FFI Python CPU callback custom call.
+inline constexpr llvm::StringRef kFFIPythonCpuCallbackCustomCallTargetName =
+    "xla_ffi_python_cpu_callback";
+
+// The target name of the Python GPU callback custom call.
+inline constexpr llvm::StringRef kPythonGpuCallbackCustomCallTargetName =
+    "xla_python_gpu_callback";
+
+// The target name of the FFI Python GPU callback custom call.
+inline constexpr llvm::StringRef kFFIPythonGpuCallbackCustomCallTargetName =
+    "xla_ffi_python_gpu_callback";
+
+// The attribute name for backend config.
+inline constexpr llvm::StringRef kXlaBackendConfigAttr = "backend_config";
+
+// Attribute name for temporarily storing the Shardy sharding during HLO
+// sdy-round-trip. It cannot match the name `kShardingAttr` ("sdy.sharding"), as
+// during sdy-round-trip, going from HLO to StableHLO, the code removes
+// attributes in the `frontend_attributes` field, making them top level. And
+// Shardy verification expects `kShardingAttr` to be of type
+// TensorShardingAttr/TensorShardingPerValueAttr - not a StringAttr.
+inline constexpr llvm::StringRef kShardingRoundTripAttr = "xla.sdy.sharding";
+
+// Attribute name for temporarily storing the Shardy sharding rule during HLO
+// sdy-round-trip. It cannot match the name `kShardingRuleAttr`
+// ("sdy.sharding_rule"), as during sdy-round-trip, going from HLO to StableHLO,
+// the code removes attributes in the `frontend_attributes` field, making them
+// top level. And Shardy verification expects `kShardingRuleAttr` to be of type
+// OpShardingRuleAttr - not a StringAttr.
+inline constexpr llvm::StringRef kShardingRuleRoundTripAttr =
+    "xla.sdy.sharding_rule";
+
+// Attribute name for temporarily storing the Shardonnay meshes during HLO
+// round-trip.
+inline constexpr llvm::StringRef kMeshesRoundTripAttr = "xla.sdy.meshes";
+
+// The target name of the custom call when round tripping during HLO
+// round-trip.
+inline constexpr llvm::StringRef kFuncResultShardingTargetName =
+    "xla.sdy.FuncResultSharding";
+
+// The target name of the ShardingGroup custom call.
+inline constexpr llvm::StringRef kShardingGroupCustomCallTargetName =
+    "xla.sdy.ShardingGroup";
+
+// Sharding group id attribute name. The attribute will be of type `int64_t`
+// and will be used to identify a group of ops that should be sharded together.
+inline constexpr llvm::StringRef kShardingGroupIdAttr =
+    "xla.sdy.sharding_group_id";
+
+// Attribute name for storing frontend attributes in XLA.
+inline constexpr llvm::StringRef kFrontendAttributesAttr =
+    "mhlo.frontend_attributes";
+
+// Attribute name for determining whether we need to import StableHLO shardings,
+// i.e., the input module doesn't contain SDY shardings as frontend attributes.
+inline constexpr llvm::StringRef kImportMhloShardings =
+    "xla.sdy.import_mhlo_shardings";
+
+// Attribute name for determining whether tuple parameters should be used for
+// the rest of the XLA pipeline.
+// TODO(b/345414638): remove this when Shardy is the first thing run in the
+// XLA pipeline, so no HLO<->MLIR round-tripping.
+inline constexpr llvm::StringRef kUseTupleArgs = "xla.sdy.use_tuple_args";
+
+// Attribute name for the in shardings of a `ManualComputationOp`.
+inline constexpr llvm::StringRef kInShardings = "xla.sdy.in_shardings";
+
+// Attribute name for the out shardings of a `ManualComputationOp`.
+inline constexpr llvm::StringRef kOutShardings = "xla.sdy.out_shardings";
+
+// Attribute name for the manual axes of a `ManualComputationOp`.
+inline constexpr llvm::StringRef kManualAxes = "xla.sdy.manual_axes";
+
+// The function name of the of the body of a `ManualComputationOp` during Shardy
+// round tripping. Used
+inline constexpr llvm::StringRef kManualComputationBodyFuncName =
+    "xla.sdy.manual_computation_body";
+
+// The target name of the custom call that changes operands from global to local
+// shape during Shardy round tripping.
+inline constexpr llvm::StringRef kGlobalToLocalShapeCallTargetName =
+    "xla.sdy.GlobalToLocalShape";
+
+// The target name of the custom call that changes results from local to global
+// shape during Shardy round tripping.
+inline constexpr llvm::StringRef kLocalToGlobalShapeCallTargetName =
+    "xla.sdy.LocalToGlobalShape";
+
+// The name of the global mesh.
+inline constexpr llvm::StringRef kGlobalMeshName = "mesh";
+
+}  //  namespace sdy
+}  //  namespace xla
+
+#endif  // SHARDY_ROUND_TRIP_IMPORT_CONSTANTS_H_