OpenXLA Project

Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

`-sdy-close-shardings`

關閉張量 sharding 並捨棄複製的軸。

`-sdy-constant-or-scalar-merger`

合併相同的常數和純量擴充，並比對分片。

對具有相同分區的常數執行輕量 CSE。

匯入管道會分割及複製常數和純量擴充，因此不會在常數子運算的不同用途之間傳播 shard。如果常數在傳播後具有相同的 shard，這個階段會合併這些常數，以節省編譯時間。詳情請參閱 -sdy-constant-or-scalar-splitter。

`-sdy-convert-global-to-local`

將 SDY 程式從全域形狀轉換為本機形狀。

根據分片屬性分割邏輯維度，將 SDY 程式從全域形狀轉換為本機形狀。

這個階段會使用類型轉換器，將 RankedTensorType 從全域邏輯形狀對應至裝置本機實體形狀。

選項

-per-dim-all-gather                     : Keep per-dimension all-gather without combining them into a single all-gather.
-combine-multi-dimension-reduce-scatter : Combine multi-dimension reduce-scatter into a single reduce-scatter.
-enable-rgv3                            : Use StableHLO ReplicaGroupV3 (mesh-axes based) for collectives.
-replica-count                          : Number of replicas (data parallelism).
-partition-count                        : Number of partitions (model parallelism).

`-sdy-drop-sharding-and-mesh`

從程式中移除網格運算和 sharding 標記。

選項

-keep-module-sharding-metadata : Whether to keep module-level sharding metadata (sdy.parameters_shardings and sdy.output_shardings).

`-sdy-drop-sharding-rules`

所有已註冊營運商的開賣活動 OpShardingRuleAttr。

`-sdy-export-named-computations`

列出 NamedComputationOp 的通話記錄。

建立一個傳遞，使用名為 NamedComputationOp 的 name 的新私人函式，將 NamedComputationOp 轉換為 CallOp。新的 FuncOp 和 CallOp 與原始 NamedComputationOp 的運算元/結果具有相同的分片。

如果模組中有與 NamedComputationOp 同名的函式，MLIR 符號表會將其變更為 {name}_#。

`-sdy-insert-explicit-reshards`

插入明確的重新分 shard，讓所有作業都有相容的分 shard。

相容的分區基本上代表作業可接受分區運算元，並產生分區結果，不需任何重新分區通訊 (請注意，作業可能仍需通訊，例如 all-reduce 或 halo-swap)。

傳播後，部分作業可能仍有不相容的分區。

請注意，如果使用軸 (或子軸) 將不對應的維度 (例如 matmul 中的非收縮維度) 分割到多個張量，或使用軸分割一個張量中的維度，但未分割另一個張量中的對應維度，則表示該運算有分割衝突。因此，在這次傳遞後，作業就不會發生衝突。

這個階段會明確插入重新分 shard 作業，讓每個作業的對應維度在所有運算元和結果中，都以相同方式分 shard，且每個軸 (或子軸) 只能用於分 shard 單一維度類型。

範例：

輸入：

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

輸出內容：

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

在上述範例中，lhs 和 rhs 都在非縮減維度的「x」軸上進行 shard，因此不相容。這個傳遞會在點號運算子之前的 rhs 上插入明確的重新分片，以便點號運算子具有相容的分片。

選項

-enable-full-version                     : Enable full version.
-mark-partial-result-with-unreduced-axes : Whether to mark partial reduction results with unreduced axes.

`-sdy-insert-func-call-reshards`

插入 func 和呼叫分片衝突的重新分片。

在結果中插入 func 和呼叫分片衝突的重新分片。

`-sdy-optimize-collectives`

最佳化集體通訊作業。

分解分割維度並融合為合併的 all-to-all，藉此消除 all-to-all 鏈結前多餘的集合排列作業。

`-sdy-pad-for-divisibility`

將具有不可分割分區的張量，填補為可分割的形狀。

`-sdy-per-instruction-partitioning`

_選擇性擷取目標指令，對這些指令執行分割器管道，產生裝置本機程式碼，然後將這些程式碼包裝在 sdy.manualcomputation 中。

為了偵錯和二分，使用獨立的 Shardy 分割器管道 (resolve-permutation-factors、reshard-to-collectives、pad-for-divisibility、convert-global-to-local) 選擇性地分割個別指令，並將產生的裝置本機程式碼包裝在 sdy.manual_computation 中。這樣一來，全域程式的其餘部分就會保持不變，因此 XLA SPMD 只會分割未包裝的全域指令。

選項

-filter                   : Filter string for selective partitioning. Can be empty (all sharded ops), comma-separated op name substrings (e.g. 'dot, pad'), or key-values like 'selectLow=0, selectHigh=10, func=subroutine'.
-enable-halo-exchange     : Implement halo exchange logic for windowed operations inside the sub-pipeline.
-replica-count            : Number of replicas (data parallelism).
-partition-count          : Number of partitions (model parallelism).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

`-sdy-propagate-to-func-results`

將 shardings 從 func 終止符傳播到 func 結果。

將函式終止符值的分片複製到對應的 func.func 結果，並保留主要函式的現有分片。

`-sdy-remove-all-gather-reduce-scatter-for-cmv1`

_移除 CMV1 的 sdy.all_gather 和 sdy.reducescatter。

移除 all-gather + dot 模式中的 all-gather。移除 dot + reduce-scatter 模式中的 reduce-scatter。這個傳遞作業是為了與集體 matmul V1 (CMV1) 相容。這是 b/432019089 的暫時解決方案。

`-sdy-remove-propagation-debug-info`

匯出時會移除傳播偵錯資訊 (傳播邊緣和來源分片)。

`-sdy-remove-sharding-groups`

Removes ShardingGroupOps after propagation.

`-sdy-remove-sub-axes-in-input-output-shardings`

移除輸入/輸出分片中的子軸。

部分 Shardy 使用者預期函式輸入/輸出內容的分區不含子軸。這個傳遞會從輸入/輸出開放維度分片中移除子軸及其尾端軸。這個階段通常會在 sdy-update-non-divisible-input-output-shardings 之後進行，確保移除子軸不會產生任何不可分割的分區。

`-sdy-reshard-to-collectives`

將 ReshardOp 轉換為各種 Shardy 集合運算。

比對重新分 shard 作業，並將其重寫為各種 Shardy 集合運算。完成後，模組中就不會再有重新分 shard 作業。

如果 keepRedundantReshards 為 true，則只會保留多餘的重新分片作業。根據預設，系統會假設已插入明確的重新分片 (sdy-insert-explicit-reshards)，且不會保留多餘的重新分片。如果可能尚未插入明確的重新分片，則應保留多餘的重新分片。

範例：

輸入：

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.reshard %arg0 <@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

輸出內容：

mesh = <"x"=2, "y"=2, "z"=2>
%0 : tensor<16x2xf32> {sdy.sharding<@mesh, \[{"x", "y", "z"}, {}\]>
%1 = sdy.all_gather \[{"y", "z"}, {}\] %arg0 out_sharding=<@mesh, \[{"x"}, {}\]> : tensor<16x2xf32>

在上述範例中，張量 %0 : tensor<16x2xf32> 會分割為 \[{"x", "y", "z"}, {}\]。接著，reshard op 會將其重新分片為 \[{"x"}, {}\]。在第一個軸上，由於重新分片後會移除後置字串 {"y", "z"}，因此我們推斷已 all-gather {"y", "z"}。第二個維度則維持不變。

選項

-keep-redundant-reshards : Whether it keeps redundant reshards or removes.

`-sdy-resolve-permutation-factors`

解決對應至 kPermutation 因素的維度資料分割問題。

使用 kPermutation 係數進行資料分割維度時，可能需要跨裝置通訊 (例如視窗的光環交換或反向的集體排列)。

如果 enableHaloExchange 為 true，則傳遞會使用可用的最佳化通訊邏輯來解析排列因素。否則，傳遞作業只會插入 sdy.reshard 作業來複製這些維度。enableHaloExchange 的預設值為 true。

選項

-enable-halo-exchange     : Implement halo exchange logic for windowed operations.
-replica-count            : Number of replicas (data parallelism).
-partition-count          : Number of partitions (model parallelism).
-rng-bit-generator-unsafe : Whether to allow unsafe partitioning of stablehlo.rng_bit_generator.

`-sdy-resolve-single-device-sharding`

將單一裝置作業降至 stablehlo.if，並由目標裝置 ID 保護。

選項

-replica-count   : Number of replicas per partition.
-partition-count : Number of partitions per replica.

`-sdy-sharding-constraint-to-reshard`

將 ShardingConstraintOp 轉換為 ReshardOp。

`-sdy-sink-data-flow-edges`

將所有 DataFlowEdgeOp 接收至輸入內容。

將每個 DataFlowEdgeOp 的分片作業移至其輸入內容 (邊緣的根目標)，並以輸入內容取代作業。

選項

-sink-debug-sharding-origins          : Whether to sink the debug sharding origins info. See `debug-sharding-origins` option in propagation for more info.
-sink-debug-propagation-edge-sharding : Whether to sink the debug propagation edge sharding info. See `debug-propagation-edge-sharding` option in propagation for more info.

`-sdy-sink-func-data-flow-edges`

將所有 FuncDataFlowEdgeOp 接收至輸入內容。

將每個 FuncDataFlowEdgeOp 的分片移至其輸入內容，並以輸入內容取代作業。

`-sdy-unflatten-call-graph`

取消扁平化呼叫圖。

取消扁平化圖表。如果函式具有相同的輸入/輸出 sharding「和」相同的來源 (如函式附加的「original_func_name」屬性所述)，則會簡化函式。

選項

-dedup-functions-fully : If true, regardless of the input and output shardings of functions, it keeps one callee function for each caller function. The default is false, meaning it will deduplicate only if the input and output shardings are the same.

`-sdy-update-non-divisible-input-output-shardings`

將 FuncOp 輸入/輸出平均分區，避免因分區無法整除而需要填補。

Shardy 使用者預期函式輸入/輸出內容可平均分割/分區，避免需要對張量進行 padding。傳播可能會導致輸入/輸出具有不可分割的資料分割，因此這個傳遞會將其更新為均勻資料分割的原始資料分割的最大維度資料分割前置字元。

`-sdy-verify-unreduced-axes`

驗證未縮減軸的使用一致性。

驗證每項運算，如果運算元的軸未縮減，運算會明確縮減軸 (例如透過 sdy.reshard)，或將軸傳遞至結果 (或為 func.call 等繫結運算)。