Skip to content

Commit

Permalink
Fix typo regarding ImportConstantsPass comment.
Browse files Browse the repository at this point in the history
The pass is still operating on `mhlo` constants not `stablehlo` constants. This is because we need to call `mhlo::createFlattenTuplePass` after which is a greedy pattern with folding.

PiperOrigin-RevId: 713982378
  • Loading branch information
bartchr808 authored and Google-ML-Automation committed Jan 10, 2025
1 parent 62e57c4 commit 8065544
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions xla/service/spmd/shardy/round_trip_common/pipeline_passes.cc
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,11 @@ void addCommonPreImportPasses(mlir::OpPassManager& pm) {
// changes happen before shardings are added to operations, to ensure the
// correct shardings are added and that they are not lost by this pass.
pm.addNestedPass<FuncOp>(mlir::mhlo::createPrepareForExportPass());
// We import `stablehlo.constant` ops to `sdy.constant` ops so that constants
// We import `mhlo.constant` ops to `sdy.constant` ops so that constants
// aren't folded in greedy pattern rewriters, which would lift them outside of
// nested regions (this undoes `WhileLoopConstantSinking` HLO pass).
// Therefore, this pass needs to be applied after any stablehlo pass that
// expects `stablehlo.constant`, and before any pass that has a greedy pattern
// Therefore, this pass needs to be applied after any MHLO pass that
// expects `mhlo.constant`, and before any pass that has a greedy pattern
// rewriter.
pm.addNestedPass<FuncOp>(createImportConstantsPass());
pm.addNestedPass<FuncOp>(mlir::mhlo::createFlattenTuplePass());
Expand Down

0 comments on commit 8065544

Please sign in to comment.