Skip to content

Commit 4434734

Browse files
abhigunjGoogle-ML-Automation
abhigunj
authored and
Google-ML-Automation
committed
[Shardy] Replace "mhlo" str with "stablehlo" from files,functions,dirs names
PiperOrigin-RevId: 722724716
1 parent 1a10f9d commit 4434734

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+278
-299
lines changed

‎xla/service/spmd/shardy/BUILD

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -38,16 +38,13 @@ cc_library(
3838
"//xla/hlo/transforms/simplifiers:hlo_dce",
3939
"//xla/hlo/transforms/simplifiers:tuple_simplifier",
4040
"//xla/hlo/translate:stablehlo",
41-
"//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
42-
"//xla/hlo/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
4341
"//xla/hlo/utils:hlo_sharding_util",
44-
"//xla/mlir_hlo:mhlo_passes",
4542
"//xla/service:computation_layout",
4643
"//xla/service:hlo_proto_cc",
4744
"//xla/service/llvm_ir:llvm_util",
48-
"//xla/service/spmd/shardy/mhlo_round_trip:mhlo_export",
49-
"//xla/service/spmd/shardy/mhlo_round_trip:mhlo_import",
5045
"//xla/service/spmd/shardy/sdy_round_trip:pipelines",
46+
"//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_export",
47+
"//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import",
5148
"//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
5249
"@com_google_absl//absl/algorithm:container",
5350
"@com_google_absl//absl/container:flat_hash_map",
@@ -86,7 +83,6 @@ cc_library(
8683
"@llvm-project//mlir:FuncExtensions",
8784
"@llvm-project//mlir:IR",
8885
"@llvm-project//mlir:Support",
89-
"@llvm-project//mlir:TensorDialect",
9086
"@shardy//shardy/dialect/sdy/ir:dialect",
9187
"@shardy//shardy/dialect/sdy/ir:register",
9288
"@stablehlo//:stablehlo_ops",
@@ -121,16 +117,7 @@ xla_cc_binary(
121117
name = "sdy_opt",
122118
srcs = ["sdy_opt_main.cc"],
123119
deps = [
124-
"//xla/mlir_hlo",
125-
"//xla/mlir_hlo:mhlo_passes",
126120
"//xla/mlir_hlo:stablehlo_extension_passes",
127-
"//xla/service/spmd/shardy/mhlo_round_trip:export_callback_custom_calls",
128-
"//xla/service/spmd/shardy/mhlo_round_trip:export_ops",
129-
"//xla/service/spmd/shardy/mhlo_round_trip:export_shardings",
130-
"//xla/service/spmd/shardy/mhlo_round_trip:mhlo_export",
131-
"//xla/service/spmd/shardy/mhlo_round_trip:mhlo_import",
132-
"//xla/service/spmd/shardy/mhlo_round_trip:shard_map_export",
133-
"//xla/service/spmd/shardy/mhlo_round_trip:shard_map_import",
134121
"//xla/service/spmd/shardy/round_trip_common:export_named_computations",
135122
"//xla/service/spmd/shardy/round_trip_common:import_backend_func_calls",
136123
"//xla/service/spmd/shardy/round_trip_common:import_constants",
@@ -144,13 +131,19 @@ xla_cc_binary(
144131
"//xla/service/spmd/shardy/sdy_round_trip:remove_size_one_axes",
145132
"//xla/service/spmd/shardy/sdy_round_trip:shard_map_export",
146133
"//xla/service/spmd/shardy/sdy_round_trip:shard_map_import",
147-
"//xla/service/spmd/shardy/sdy_round_trip/test_utils:mhlo_to_hlo_to_mhlo",
134+
"//xla/service/spmd/shardy/sdy_round_trip/test_utils:stablehlo_to_hlo_to_stablehlo",
148135
"//xla/service/spmd/shardy/sdy_round_trip/test_utils:testing_pipeline",
136+
"//xla/service/spmd/shardy/stablehlo_round_trip:export_callback_custom_calls",
137+
"//xla/service/spmd/shardy/stablehlo_round_trip:export_ops",
138+
"//xla/service/spmd/shardy/stablehlo_round_trip:export_shardings",
139+
"//xla/service/spmd/shardy/stablehlo_round_trip:shard_map_export",
140+
"//xla/service/spmd/shardy/stablehlo_round_trip:shard_map_import",
141+
"//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_export",
142+
"//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import",
149143
"@llvm-project//mlir:AllPassesAndDialects",
150144
"@llvm-project//mlir:FuncExtensions",
151145
"@llvm-project//mlir:IR",
152146
"@llvm-project//mlir:MlirOptLib",
153-
"@shardy//shardy/dialect/sdy/ir:dialect",
154147
"@shardy//shardy/dialect/sdy/ir:register",
155148
"@shardy//shardy/dialect/sdy/transforms:passes",
156149
],

‎xla/service/spmd/shardy/constants.h

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ limitations under the License.
2121
namespace xla {
2222
namespace sdy {
2323

24-
// The attribute name for attributes in MHLO ops.
24+
// TODO(b/393406516): [Shardy] Change Attributes names from *mhlo* to
25+
// *stablehlo*
2526
inline constexpr llvm::StringRef kMhloAttributesAttr = "mhlo.attributes";
2627

2728
// The attribute name for xla::HloSharding.
@@ -59,17 +60,17 @@ inline constexpr llvm::StringRef kXlaBackendConfigAttr = "backend_config";
5960

6061
// Attribute name for temporarily storing the Shardy sharding during HLO
6162
// sdy-round-trip. It cannot match the name `kShardingAttr` ("sdy.sharding"), as
62-
// during sdy-round-trip, going from HLO to MHLO, the code removes attributes
63-
// in the `frontend_attributes` field, making them top level. And Shardy
64-
// verification expects `kShardingAttr` to be of type
63+
// during sdy-round-trip, going from HLO to StableHLO, the code removes
64+
// attributes in the `frontend_attributes` field, making them top level. And
65+
// Shardy verification expects `kShardingAttr` to be of type
6566
// TensorShardingAttr/TensorShardingPerValueAttr - not a StringAttr.
6667
inline constexpr llvm::StringRef kShardingRoundTripAttr = "xla.sdy.sharding";
6768

6869
// Attribute name for temporarily storing the Shardy sharding rule during HLO
6970
// sdy-round-trip. It cannot match the name `kShardingRuleAttr`
70-
// ("sdy.sharding_rule"), as during sdy-round-trip, going from HLO to MHLO, the
71-
// code removes attributes in the `frontend_attributes` field, making them top
72-
// level. And Shardy verification expects `kShardingRuleAttr` to be of type
71+
// ("sdy.sharding_rule"), as during sdy-round-trip, going from HLO to StableHLO,
72+
// the code removes attributes in the `frontend_attributes` field, making them
73+
// top level. And Shardy verification expects `kShardingRuleAttr` to be of type
7374
// OpShardingRuleAttr - not a StringAttr.
7475
inline constexpr llvm::StringRef kShardingRuleRoundTripAttr =
7576
"xla.sdy.sharding_rule";
@@ -96,7 +97,7 @@ inline constexpr llvm::StringRef kShardingGroupIdAttr =
9697
inline constexpr llvm::StringRef kFrontendAttributesAttr =
9798
"mhlo.frontend_attributes";
9899

99-
// Attribute name for determining whether we need to import MHLO shardings,
100+
// Attribute name for determining whether we need to import StableHLO shardings,
100101
// i.e., the input module doesn't contain SDY shardings as frontend attributes.
101102
inline constexpr llvm::StringRef kImportMhloShardings =
102103
"xla.sdy.import_mhlo_shardings";

‎xla/service/spmd/shardy/round_trip_common/BUILD

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,6 @@ cc_library(
107107
":import_constants",
108108
":import_sdy_custom_calls",
109109
":open_while_free_vars_sharding",
110-
"//xla/mlir_hlo:mhlo_passes",
111110
"//xla/mlir_hlo:stablehlo_extension_passes",
112111
"@llvm-project//mlir:FuncDialect",
113112
"@llvm-project//mlir:Pass",

‎xla/service/spmd/shardy/round_trip_common/pipeline_passes.cc

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ limitations under the License.
2020
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
2121
#include "mlir/Transforms/Passes.h"
2222
#include "stablehlo/transforms/Passes.h"
23-
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
2423
#include "xla/mlir_hlo/stablehlo_ext/transforms/passes.h"
2524
#include "xla/service/spmd/shardy/round_trip_common/import_backend_func_calls.h"
2625
#include "xla/service/spmd/shardy/round_trip_common/import_constants.h"

‎xla/service/spmd/shardy/round_trip_common/pipeline_passes.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,17 @@ limitations under the License.
2121
namespace xla {
2222
namespace sdy {
2323

24-
// Adds the common import passes for both the SDY and MHLO import
24+
// Adds the common import passes for both the SDY and StableHLO import
2525
// pipelines that need to be called before each pipeline converts an HLO
2626
// sharding/SDY sharding string into an `sdy.sharding` attribute.
2727
void addCommonPreImportPasses(mlir::OpPassManager& pm);
2828

29-
// Adds the common import passes for both the SDY and MHLO import
29+
// Adds the common import passes for both the SDY and StableHLO import
3030
// pipelines that need to be called after each pipeline converts an HLO
3131
// sharding/SDY sharding string into an `sdy.sharding` attribute.
3232
void addCommonPostImportPasses(mlir::OpPassManager& pm);
3333

34-
// Adds the common export passes for both the SDY and MHLO import
34+
// Adds the common export passes for both the SDY and StableHLO import
3535
// pipelines that need to be called before each pipeline converts an HLO
3636
// sharding/SDY sharding string into an `sdy.sharding` attribute.
3737
void addCommonPreExportPasses(mlir::OpPassManager& pm);

‎xla/service/spmd/shardy/sdy_opt_main.cc

Lines changed: 18 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,9 @@ limitations under the License.
1717
#include "mlir/IR/DialectRegistry.h"
1818
#include "mlir/InitAllPasses.h"
1919
#include "mlir/Tools/mlir-opt/MlirOptMain.h"
20-
#include "shardy/dialect/sdy/ir/dialect.h"
2120
#include "shardy/dialect/sdy/ir/register.h"
2221
#include "shardy/dialect/sdy/transforms/passes.h"
23-
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
24-
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
25-
#include "xla/service/spmd/shardy/mhlo_round_trip/export_callback_custom_calls.h"
26-
#include "xla/service/spmd/shardy/mhlo_round_trip/export_ops.h"
27-
#include "xla/service/spmd/shardy/mhlo_round_trip/export_shardings.h"
28-
#include "xla/service/spmd/shardy/mhlo_round_trip/mhlo_export.h"
29-
#include "xla/service/spmd/shardy/mhlo_round_trip/mhlo_import.h"
30-
#include "xla/service/spmd/shardy/mhlo_round_trip/shard_map_export.h"
31-
#include "xla/service/spmd/shardy/mhlo_round_trip/shard_map_import.h"
22+
#include "xla/mlir_hlo/stablehlo_ext/transforms/passes.h"
3223
#include "xla/service/spmd/shardy/round_trip_common/export_named_computations.h"
3324
#include "xla/service/spmd/shardy/round_trip_common/import_backend_func_calls.h"
3425
#include "xla/service/spmd/shardy/round_trip_common/import_constants.h"
@@ -42,37 +33,42 @@ limitations under the License.
4233
#include "xla/service/spmd/shardy/sdy_round_trip/remove_size_one_axes.h"
4334
#include "xla/service/spmd/shardy/sdy_round_trip/shard_map_export.h"
4435
#include "xla/service/spmd/shardy/sdy_round_trip/shard_map_import.h"
45-
#include "xla/service/spmd/shardy/sdy_round_trip/test_utils/mhlo_to_hlo_to_mhlo.h"
36+
#include "xla/service/spmd/shardy/sdy_round_trip/test_utils/stablehlo_to_hlo_to_stablehlo.h"
4637
#include "xla/service/spmd/shardy/sdy_round_trip/test_utils/testing_pipeline.h"
38+
#include "xla/service/spmd/shardy/stablehlo_round_trip/export_callback_custom_calls.h"
39+
#include "xla/service/spmd/shardy/stablehlo_round_trip/export_ops.h"
40+
#include "xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.h"
41+
#include "xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.h"
42+
#include "xla/service/spmd/shardy/stablehlo_round_trip/shard_map_import.h"
43+
#include "xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_export.h"
44+
#include "xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h"
4745

4846
int main(int argc, char** argv) {
4947
mlir::registerAllPasses();
50-
mlir::mhlo::registerAllMhloPasses();
51-
48+
mlir::stablehlo_ext::registerPasses();
5249
mlir::DialectRegistry dialects;
5350
mlir::sdy::registerAllDialects(dialects);
54-
dialects.insert<mlir::mhlo::MhloDialect>();
5551
mlir::func::registerAllExtensions(dialects);
5652

5753
// Register all SDY passes and pipelines.
5854
mlir::sdy::registerAllSdyPassesAndPipelines();
5955

60-
xla::sdy::registerMhloImportPipeline();
61-
xla::sdy::registerMhloImportShardingsPass();
62-
xla::sdy::registerMhloRoundTripShardMapImportPass();
56+
xla::sdy::registerStablehloImportPipeline();
57+
xla::sdy::registerStablehloImportShardingsPass();
58+
xla::sdy::registerStablehloRoundTripShardMapImportPass();
6359
xla::sdy::registerImportSdyCustomCallsPass();
6460
xla::sdy::registerOpenWhileFreeVarsShardingPass();
6561
xla::sdy::registerImportBackendFuncCallsPass();
6662
xla::sdy::registerImportConstantsPass();
6763

68-
xla::sdy::registerMhloExportPipeline();
69-
xla::sdy::registerMhloExportShardingsPass();
70-
xla::sdy::registerMhloRoundTripExportCallbackCustomCallsPass();
71-
xla::sdy::registerMhloRoundTripShardMapExportPass();
64+
xla::sdy::registerStablehloExportPipeline();
65+
xla::sdy::registerStablehloExportShardingsPass();
66+
xla::sdy::registerStablehloRoundTripExportCallbackCustomCallsPass();
67+
xla::sdy::registerStablehloRoundTripShardMapExportPass();
7268
xla::sdy::registerExportNamedComputationsPass();
7369
xla::sdy::registerExportOpsPass();
7470

75-
xla::sdy::registerSdyRoundTripMhloToHloToMhloPass();
71+
xla::sdy::registerSdyRoundTripStablehloToHloToStablehloPass();
7672
xla::sdy::registerSdyRoundTripExportShardyAttrsPass();
7773
xla::sdy::registerSdyRoundTripImportCallbackCustomCallsPass();
7874
xla::sdy::registerSdyRoundTripImportShardyAttrsPass();

‎xla/service/spmd/shardy/sdy_round_trip/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,9 +155,9 @@ cc_library(
155155
":shard_map_export",
156156
":shard_map_import",
157157
"//xla/service:hlo_proto_cc",
158-
"//xla/service/spmd/shardy/mhlo_round_trip:export_shardings",
159158
"//xla/service/spmd/shardy/round_trip_common:export_named_computations",
160159
"//xla/service/spmd/shardy/round_trip_common:pipeline_passes",
160+
"//xla/service/spmd/shardy/stablehlo_round_trip:export_shardings",
161161
"@llvm-project//mlir:Pass",
162162
"@llvm-project//mlir:Support",
163163
],

‎xla/service/spmd/shardy/sdy_round_trip/export_ops.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ class SdyRoundTripExportOpsPass
153153
}
154154

155155
StringRef getDescription() const override {
156-
return "Exports Shardonnay ops to MHLO ops.";
156+
return "Exports Shardonnay ops to StableHLO ops.";
157157
}
158158
};
159159

‎xla/service/spmd/shardy/sdy_round_trip/export_ops.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ limitations under the License.
2323
namespace xla {
2424
namespace sdy {
2525

26-
// Creates a pass that converts Shardy ops to MHLO ops (except
26+
// Creates a pass that converts Shardy ops to StableHLO ops (except
2727
// sdy::ManualComputationOp).
2828
std::unique_ptr<mlir::Pass> createSdyRoundTripExportOpsPass();
2929

‎xla/service/spmd/shardy/sdy_round_trip/export_shardy_attrs.cc

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -97,9 +97,9 @@ LogicalResult exportFunc(FuncOp funcOp, OpBuilder& builder) {
9797
for (mlir::OpOperand& returnOperand : terminatorOp->getOpOperands()) {
9898
if (auto sharding = funcOp.getResultAttrOfType<TensorShardingAttr>(
9999
returnOperand.getOperandNumber(), kShardingAttr)) {
100-
// We cannot save the result shardings as frontend attributes. MHLO->HLO
101-
// conversion converts `mhlo.sharding`s on the results to a tuple
102-
// sharding on the ROOT instruction, but it discards the frontend
100+
// We cannot save the result shardings as frontend attributes.
101+
// StableHLO->HLO conversion converts `mhlo.sharding`s on the results to a
102+
// tuple sharding on the ROOT instruction, but it discards the frontend
103103
// attributes on the MLIR results. Thus, instead of making the conversion
104104
// handle that, push the shardings to a temporary custom op for each
105105
// FuncOp result with a sharding. Then during import, we will copy the
@@ -152,15 +152,16 @@ class SdyRoundTripExportShardyAttrsPass
152152
}
153153
}
154154

155-
SmallVector<NamedAttribute> mhloMeshes;
156-
// Saves the MeshOps for MHLO<->HLO round-trip and removes them from the
157-
// ModuleOp.
155+
SmallVector<NamedAttribute> stablehloMeshes;
156+
// Saves the MeshOps for StableHLO<->HLO round-trip and removes them from
157+
// the ModuleOp.
158158
for (MeshOp meshOp : moduleOp.getOps<MeshOp>()) {
159-
mhloMeshes.emplace_back(meshOp.getSymNameAttr(), meshOp.getMeshAttr());
159+
stablehloMeshes.emplace_back(meshOp.getSymNameAttr(),
160+
meshOp.getMeshAttr());
160161
}
161-
if (!mhloMeshes.empty()) {
162+
if (!stablehloMeshes.empty()) {
162163
setFrontendAttribute(moduleOp, kMeshesRoundTripAttr,
163-
DictionaryAttr::get(context, mhloMeshes));
164+
DictionaryAttr::get(context, stablehloMeshes));
164165
}
165166
}
166167

‎xla/service/spmd/shardy/sdy_round_trip/pipelines.cc

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ limitations under the License.
2121
#include "mlir/Pass/PassRegistry.h"
2222
#include "mlir/Support/LLVM.h"
2323
#include "xla/service/hlo.pb.h"
24-
#include "xla/service/spmd/shardy/mhlo_round_trip/export_shardings.h"
2524
#include "xla/service/spmd/shardy/round_trip_common/export_named_computations.h"
2625
#include "xla/service/spmd/shardy/round_trip_common/pipeline_passes.h"
2726
#include "xla/service/spmd/shardy/sdy_round_trip/export_ops.h"
@@ -31,6 +30,7 @@ limitations under the License.
3130
#include "xla/service/spmd/shardy/sdy_round_trip/remove_size_one_axes.h"
3231
#include "xla/service/spmd/shardy/sdy_round_trip/shard_map_export.h"
3332
#include "xla/service/spmd/shardy/sdy_round_trip/shard_map_import.h"
33+
#include "xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.h"
3434

3535
namespace xla {
3636
namespace sdy {
@@ -41,11 +41,11 @@ void addSdyRoundTripExportPipeline(mlir::OpPassManager& pm) {
4141
pm.addPass(createExportNamedComputationsPass());
4242
pm.addPass(createSdyRoundTripExportOpsPass());
4343
pm.addPass(createSdyRoundTripShardMapExportPass());
44-
// Preserve the SDY shardings for `createExportMhloShardingsPass` so that
44+
// Preserve the SDY shardings for `createExportStablehloShardingsPass` so that
4545
// we have both `mhlo.sharding`s and hidden `sdy.sharding`s on the module. We
4646
// want to have `mhlo.sharding`s for Pathways to read from.
4747
pm.addPass(createSdyRoundTripExportShardyAttrsPass());
48-
pm.addPass(createExportMhloShardingsPass());
48+
pm.addPass(createExportStablehloShardingsPass());
4949
}
5050

5151
void addSdyRoundTripImportPipeline(mlir::OpPassManager& pm) {
@@ -60,15 +60,15 @@ void addSdyRoundTripImportPipeline(mlir::OpPassManager& pm) {
6060
void registerSdyRoundTripExportPipeline() {
6161
PassPipelineRegistration<> exportPipeline(
6262
"xla-sdy-round-trip-export-pipeline",
63-
"Run passes to export the SDY (Shardy) dialect into an MHLO module, "
63+
"Run passes to export the SDY (Shardy) dialect into an StableHLO module, "
6464
"but with the SDY ops/attrs saved for roundtripping.",
6565
addSdyRoundTripExportPipeline);
6666
}
6767

6868
void registerSdyRoundTripImportPipeline() {
6969
PassPipelineRegistration<> importPipeline(
7070
"xla-sdy-round-trip-import-pipeline",
71-
"Run passes to import an mhlo module into the SDY (Shardy) dialect.",
71+
"Run passes to import an StableHLO module into the SDY (Shardy) dialect.",
7272
addSdyRoundTripImportPipeline);
7373
}
7474

‎xla/service/spmd/shardy/sdy_round_trip/pipelines.h

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ namespace xla {
2222
namespace sdy {
2323

2424
// Add the xla-sdy-round-trip-export-pipeline in `pm`. The pipeline,
25-
// including a sequence of passes, exports the Shardy dialect into an MHLO
25+
// including a sequence of passes, exports the Shardy dialect into an StableHLO
2626
// module with no XLA shardings, but SDY shardings and meshes saved as string
2727
// frontend attributes.
2828
//
@@ -34,8 +34,9 @@ namespace sdy {
3434
// MLIR (see b/335666088). So this pass will eventually be removed.
3535
void addSdyRoundTripExportPipeline(mlir::OpPassManager& pm);
3636

37-
// Add the xla-sdy-round-trip-import-pipeline in `pm`. The pipeline, including a
38-
// sequence of passes, imports an MHLO module into the SDY (Shardy) dialect.
37+
// Add the xla-sdy-round-trip-import-pipeline in `pm`. The pipeline,
38+
// including a sequence of passes, imports an StableHLO module into the
39+
// SDY (Shardy) dialect.
3940
//
4041
// The module is assumed to have `kShardingRoundTripAttr` and
4142
// `kMeshesRoundTripAttr`.
@@ -48,8 +49,8 @@ void registerSdyRoundTripExportPipeline();
4849
void registerSdyRoundTripImportPipeline();
4950

5051
// Register the xla-sdy-round-trip-testing-pipeline.
51-
// This takes an SDY module, exports it to MHLO while saving the SDY attrs
52-
// and meshes, goes to HLO, back to MHLO, and then back to SDY.
52+
// This takes an SDY module, exports it to StableHLO while saving the SDY attrs
53+
// and meshes, goes to HLO, back to StableHLO, and then back to SDY.
5354
// This is for testing roundtripping SDY modules, but should be eventually
5455
// removed as part of b/335666088.
5556
void registerSdyRoundTripTestingPipeline();

‎xla/service/spmd/shardy/sdy_round_trip/test_utils/BUILD

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,12 @@ package_group(
1616
)
1717

1818
cc_library(
19-
name = "mhlo_to_hlo_to_mhlo",
20-
srcs = ["mhlo_to_hlo_to_mhlo.cc"],
21-
hdrs = ["mhlo_to_hlo_to_mhlo.h"],
19+
name = "stablehlo_to_hlo_to_stablehlo",
20+
srcs = ["stablehlo_to_hlo_to_stablehlo.cc"],
21+
hdrs = ["stablehlo_to_hlo_to_stablehlo.h"],
2222
deps = [
2323
"//xla/hlo/ir:hlo",
2424
"//xla/hlo/translate:stablehlo",
25-
"//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
26-
"//xla/mlir_hlo:mhlo_passes",
2725
"//xla/service:hlo_module_config",
2826
"//xla/service:hlo_proto_cc",
2927
"@com_google_absl//absl/log",
@@ -43,7 +41,7 @@ cc_library(
4341
srcs = ["testing_pipeline.cc"],
4442
hdrs = ["testing_pipeline.h"],
4543
deps = [
46-
":mhlo_to_hlo_to_mhlo",
44+
":stablehlo_to_hlo_to_stablehlo",
4745
"//xla/service/spmd/shardy/sdy_round_trip:pipelines",
4846
"@llvm-project//mlir:Pass",
4947
],

0 commit comments

Comments
 (0)