From aab71ea7014f58bfafe865b14918ffa069fb0856 Mon Sep 17 00:00:00 2001 From: Bill Varcho Date: Fri, 4 Sep 2026 14:43:19 -0700 Subject: [PATCH 01/12] [SDY][re-land] Support subaxes in ReplicaGroupMeshAxesAttr Reverts 36a0e001b825227b3113163384fb8d4b31a802c2 PiperOrigin-RevId: 976503543 --- .../xla/third_party/stablehlo/temporary.patch | 591 ++++++++++++++++++ 1 file changed, 591 insertions(+) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index b1e2e8b7ff4ea0..fcff1f3810b3a6 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -15,4 +15,595 @@ # # This file is automatically generated by generate_patch tool. # Do not edit directly. +diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +--- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ++++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +@@ -20,6 +20,7 @@ + #include + + #include "llvm/ADT/DenseSet.h" ++#include "llvm/ADT/STLExtras.h" + #include "llvm/ADT/SmallVector.h" + #include "llvm/ADT/StringRef.h" + #include "mlir/IR/Attributes.h" +@@ -31,57 +32,53 @@ + namespace mlir { + namespace stablehlo { + +-static SmallVector> ++namespace { ++ ++struct ReindexedAxes { ++ SmallVector splitAxisSizes; ++ SmallVector groupedAxisIndices; ++}; ++ ++// Generates replica groups from the reshaped mesh axis sizes and the indices of ++// the communication axes using a Reshape-Transpose permutation. ++SmallVector> + flattenedReplicaGroupsFromTransposePermutation( +- const SmallVector& meshAxisNames, +- const SmallVector& commAxisNames, +- const llvm::DenseSet& commAxisSet, +- const SmallVector& axisSizes, +- const SmallVector& deviceIds, int64_t totalDevices) { +- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. ++ ArrayRef axisSizes, ArrayRef groupedAxisIndices, ++ ArrayRef deviceIds, int64_t totalDevices) { ++ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), ++ groupedAxisIndices.end()); + SmallVector transposeAxes; + // Non-grouped axes first +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (!commAxisSet.count(meshAxisNames[i])) { ++ for (size_t i = 0; i < axisSizes.size(); ++i) { ++ if (!groupedAxisSet.count(i)) { + transposeAxes.push_back(i); + } + } +- // Grouped axes +- for (const auto& name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- transposeAxes.push_back(i); +- break; +- } +- } +- } +- +- SmallVector transposedSizes(meshAxisNames.size()); +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { ++ // Grouped axes in the specified order ++ for (int64_t idx : groupedAxisIndices) { ++ transposeAxes.push_back(idx); ++ } ++ ++ SmallVector transposedSizes(axisSizes.size()); ++ for (size_t i = 0; i < axisSizes.size(); ++i) { + transposedSizes[i] = axisSizes[transposeAxes[i]]; + } + +- // Compute strides for original shape +- SmallVector originalStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ // Compute strides for reshaped shape ++ SmallVector originalStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; + } + + // Compute strides for transposed shape +- SmallVector transposedStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ SmallVector transposedStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; + } + +- // Generate chunks + int64_t numDevicesPerGroup = 1; +- for (auto name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- numDevicesPerGroup *= axisSizes[i]; +- break; +- } +- } ++ for (int64_t idx : groupedAxisIndices) { ++ numDevicesPerGroup *= axisSizes[idx]; + } + int64_t numGroups = totalDevices / numDevicesPerGroup; + +@@ -93,7 +90,7 @@ + for (int64_t j = 0; j < numDevicesPerGroup; ++j) { + int64_t linearTransposeIdx = i * numDevicesPerGroup + j; + int64_t originalIndex = 0; +- for (size_t k = 0; k < meshAxisNames.size(); ++k) { ++ for (size_t k = 0; k < axisSizes.size(); ++k) { + int64_t coord = + (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; + originalIndex += coord * originalStrides[transposeAxes[k]]; +@@ -102,9 +99,135 @@ + } + groups.push_back(std::move(group)); + } +- + return groups; + } ++ ++// Splits mesh axes based on sub-axis references and computes the corresponding ++// indices for the communication axes. ++FailureOr computeReindexedAxes(ArrayRef axesInMesh, ++ ArrayAttr commAxes, ++ Location loc) { ++ ReindexedAxes result; ++ ++ // Validate commAxes and verify that all mesh axes exist and have valid sizes. ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::dyn_cast(attr); ++ if (!shloAxisRef) { ++ return emitError(loc) << "expected AxisRefAttr in comm_axes"; ++ } ++ StringRef axisName = shloAxisRef.getName(); ++ bool found = false; ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ found = true; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ int64_t preSize = subAxisInfo.getPreSize(); ++ int64_t size = subAxisInfo.getSize(); ++ if (preSize < 1 || size < 1) { ++ return emitError(loc) ++ << "sub-axis pre_size and size must be at least 1"; ++ } ++ int64_t nextPreSize = preSize * size; ++ if (nextPreSize > meshAxis.getSize() || ++ meshAxis.getSize() % nextPreSize != 0) { ++ return emitError(loc) ++ << "sub-axis (pre_size * size) must divide mesh axis size"; ++ } ++ } ++ break; ++ } ++ } ++ if (!found) { ++ return emitError(loc) ++ << "axis '" << axisName << "' not found in mesh definition"; ++ } ++ } ++ ++ // Split each mesh axis according to the referenced subaxes. ++ struct SplitDim { ++ StringRef axisName; ++ int64_t preSize; ++ int64_t size; ++ int64_t dimIndex; ++ }; ++ SmallVector splitDims; ++ ++ for (auto meshAxis : axesInMesh) { ++ StringRef axisName = meshAxis.getName(); ++ int64_t axisSize = meshAxis.getSize(); ++ ++ if (axisSize == 1) { ++ int64_t dimIdx = result.splitAxisSizes.size(); ++ result.splitAxisSizes.push_back(1); ++ splitDims.push_back({axisName, /*preSize=*/1, /*size=*/1, dimIdx}); ++ continue; ++ } ++ ++ SmallVector preSizes = {1, axisSize}; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ if (shloAxisRef.getName() == axisName) { ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ preSizes.push_back(subAxisInfo.getPreSize()); ++ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); ++ } ++ } ++ } ++ ++ llvm::sort(preSizes); ++ preSizes.erase(llvm::unique(preSizes), preSizes.end()); ++ ++ for (size_t j = 0; j < preSizes.size() - 1; ++j) { ++ int64_t segPreSize = preSizes[j]; ++ int64_t segSize = preSizes[j + 1] / segPreSize; ++ int64_t dimIdx = result.splitAxisSizes.size(); ++ result.splitAxisSizes.push_back(segSize); ++ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); ++ } ++ } ++ ++ // Map each communication axis to its corresponding split dimension. ++ llvm::DenseSet groupedSet; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ StringRef axisName = shloAxisRef.getName(); ++ int64_t reqPreSize = 1; ++ int64_t reqSize = 0; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ reqPreSize = subAxisInfo.getPreSize(); ++ reqSize = subAxisInfo.getSize(); ++ } else { ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ reqSize = meshAxis.getSize(); ++ break; ++ } ++ } ++ } ++ ++ bool matched = false; ++ for (const auto& splitDim : splitDims) { ++ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && ++ splitDim.size == reqSize) { ++ if (!groupedSet.insert(splitDim.dimIndex).second) { ++ return emitError(loc) ++ << "Duplicate or overlapping communication axis: " << axisName; ++ } ++ result.groupedAxisIndices.push_back(splitDim.dimIndex); ++ matched = true; ++ break; ++ } ++ } ++ if (!matched) { ++ return emitError(loc) << "Invalid or overlapping communication axis on '" ++ << axisName << "'"; ++ } ++ } ++ ++ return result; ++} ++ ++} // namespace + + FailureOr>> flattenReplicaGroupMeshAxes( + Attribute meshAttr, ArrayAttr commAxes, Location loc) { +@@ -120,34 +243,13 @@ + if (!mesh) + return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); + +- auto axesInMesh = mesh.getAxes(); +- +- // Identify which axes are communication axes. +- llvm::SmallVector commAxisNames; +- llvm::DenseSet commAxisSet; +- for (auto attr : commAxes) { +- auto shloAxisRef = llvm::dyn_cast(attr); +- if (!shloAxisRef) { +- return emitError(loc) << "expected AxisRefAttr in comm_axes"; +- } +- if (shloAxisRef.getSubAxisInfo()) { +- return emitError(loc) << "Subaxes are not supported in " +- "flattenReplicaGroupMeshAxes"; +- } +- commAxisNames.push_back(shloAxisRef.getName()); +- commAxisSet.insert(shloAxisRef.getName()); +- } +- +- // Calculate total devices and axis sizes ++ FailureOr reindexedAxes = ++ computeReindexedAxes(mesh.getAxes(), commAxes, loc); ++ if (failed(reindexedAxes)) return failure(); + + int64_t totalDevices = 1; +- SmallVector axisSizes; +- SmallVector meshAxisNames; +- for (auto meshAxis : axesInMesh) { +- auto typedMeshAxis = llvm::cast(meshAxis); +- axisSizes.push_back(typedMeshAxis.getSize()); +- meshAxisNames.push_back(typedMeshAxis.getName()); +- totalDevices *= typedMeshAxis.getSize(); ++ for (auto meshAxis : mesh.getAxes()) { ++ totalDevices *= llvm::cast(meshAxis).getSize(); + } + + SmallVector deviceIds; +@@ -160,8 +262,8 @@ + } + + return flattenedReplicaGroupsFromTransposePermutation( +- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, +- totalDevices); ++ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, ++ deviceIds, totalDevices); + } + + } // namespace stablehlo +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir +--- stablehlo/stablehlo/tests/interpret/all_gather.mlir ++++ stablehlo/stablehlo/tests/interpret/all_gather.mlir +@@ -133,3 +133,41 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { ++ %result = "stablehlo.all_gather"(%operand) { ++ all_gather_dim = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<1xi64>) -> tensor<2xi64> ++ return %result : tensor<2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], ++ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) ++ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> ++ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir +--- stablehlo/stablehlo/tests/interpret/all_reduce.mlir ++++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir +@@ -135,3 +135,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.all_reduce"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, ++ channel_handle = #stablehlo.channel_handle ++ } : (tensor<1xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], ++ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir +--- stablehlo/stablehlo/tests/interpret/all_to_all.mlir ++++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir +@@ -172,3 +172,43 @@ + func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { ++ %result = "stablehlo.all_to_all"(%operand) { ++ split_dimension = 0 : i64, ++ concat_dimension = 1 : i64, ++ split_count = 2 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2x1xi64>) -> tensor<1x2xi64> ++ return %result : tensor<1x2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> ++ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> ++ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> ++ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> ++ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> ++ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> ++ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> ++ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], ++ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] ++ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, ++ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> ++ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, ++ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) ++ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +--- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ++++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +@@ -90,3 +90,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.reduce_scatter"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ scatter_dimension = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> ++ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> ++ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> ++ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> ++ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> ++ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> ++ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> ++ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], ++ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] ++ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +--- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir ++++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +@@ -5,7 +5,7 @@ + + // CHECK-LABEL: @all_reduce_rgv3 + func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -19,7 +19,7 @@ + + // CHECK-LABEL: @all_gather_rgv3 + func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_gather"(%arg0) { + all_gather_dim = 0 : i64, + replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> +@@ -30,7 +30,7 @@ + + // CHECK-LABEL: @all_to_all_rgv3 + func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_to_all"(%arg0) { + concat_dimension = 0 : i64, + split_dimension = 0 : i64, +@@ -44,7 +44,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh + func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -59,7 +59,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh_dev + func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -84,4 +84,56 @@ + } : (tensor<4xf32>) -> tensor<4xf32> + return %0 : tensor<4xf32> + } ++ ++ // CHECK-LABEL: @all_reduce_subaxis ++ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_1 ++ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_2 ++ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_size_one_axis ++ func.func @all_reduce_size_one_axis(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } + } From 449fef9d07fdf30f5cd24d39396b7143fd18b99d Mon Sep 17 00:00:00 2001 From: Daniel Sosa Date: Fri, 4 Sep 2026 14:48:42 -0700 Subject: [PATCH 02/12] Reverts changelist 976362623 PiperOrigin-RevId: 976505695 --- ci/official/envs/linux_x86_cuda | 2 +- ci/official/envs/linux_x86_cuda13_nvcc | 2 +- .../gpu_build/parallel_gpu_execute.sh | 45 ++++++++----------- .../gpu_build/parallel_gpu_execute.sh | 45 ++++++++----------- 4 files changed, 40 insertions(+), 54 deletions(-) diff --git a/ci/official/envs/linux_x86_cuda b/ci/official/envs/linux_x86_cuda index 5e6ce4d7bef58d..dd93e2f4ded292 100644 --- a/ci/official/envs/linux_x86_cuda +++ b/ci/official/envs/linux_x86_cuda @@ -14,7 +14,7 @@ # ============================================================================== source ci/official/envs/linux_x86 export TF_FORCE_GPU_ALLOW_GROWTH=true -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true --local_test_jobs=16" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow" diff --git a/ci/official/envs/linux_x86_cuda13_nvcc b/ci/official/envs/linux_x86_cuda13_nvcc index bc8fc063157615..a8306754723f4f 100644 --- a/ci/official/envs/linux_x86_cuda13_nvcc +++ b/ci/official/envs/linux_x86_cuda13_nvcc @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== source ci/official/envs/linux_x86 -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version --local_test_jobs=16" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda_13_nvcc TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow_cuda13" diff --git a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh index 137897dd7c99c2..a00dcbc3f3404a 100755 --- a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -53,38 +53,31 @@ TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})" shift # ******************************************************************* -LOCK_DIR="${TF_LOCK_DIR:-/var/lock}" -mkdir -p "$LOCK_DIR" +mkdir -p /var/lock # Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU # slots to run a test at. # # Prefer to allocate 1 test per GPU over 4 tests on 1 GPU. # So, we iterate over TF_TESTS_PER_GPU first. -MAX_ATTEMPTS=30 -for attempt in $(seq 1 $MAX_ATTEMPTS); do - for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do - for i in `seq 0 $((TF_GPU_COUNT-1))`; do - exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exec {lock_fd}>&- - exit $return_code - fi - exec {lock_fd}>&- - done +for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exit $return_code + fi done - sleep 1 done -echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on, exiting with failure..." exit 1 diff --git a/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh index 137897dd7c99c2..a00dcbc3f3404a 100755 --- a/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/third_party/xla/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -53,38 +53,31 @@ TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})" shift # ******************************************************************* -LOCK_DIR="${TF_LOCK_DIR:-/var/lock}" -mkdir -p "$LOCK_DIR" +mkdir -p /var/lock # Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU # slots to run a test at. # # Prefer to allocate 1 test per GPU over 4 tests on 1 GPU. # So, we iterate over TF_TESTS_PER_GPU first. -MAX_ATTEMPTS=30 -for attempt in $(seq 1 $MAX_ATTEMPTS); do - for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do - for i in `seq 0 $((TF_GPU_COUNT-1))`; do - exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exec {lock_fd}>&- - exit $return_code - fi - exec {lock_fd}>&- - done +for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exit $return_code + fi done - sleep 1 done -echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on, exiting with failure..." exit 1 From a09768aec8dfe9082675b9028a6528d8e63f9c15 Mon Sep 17 00:00:00 2001 From: Bhatu Date: Fri, 4 Sep 2026 15:37:42 -0700 Subject: [PATCH 03/12] Implement constraint propagation for kExpm1 in ConstraintPropagator. PiperOrigin-RevId: 976526897 --- .../xla/xla/tests/constraint_propagator.cc | 36 ++++++++- .../xla/xla/tests/constraint_propagator.h | 2 + .../xla/tests/constraint_propagator_test.cc | 78 +++++++++++++++++++ 3 files changed, 114 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index ac980eb7a61ec1..bb0b2f2448d70c 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -440,7 +440,8 @@ void ConstraintPropagator::ComputeMaxAddReductionElementsPerExp( consumer_add_reduction_elements = it->second; } - if (instruction->opcode() == HloOpcode::kExp) { + if (instruction->opcode() == HloOpcode::kExp || + instruction->opcode() == HloOpcode::kExpm1) { max_add_reduction_elements_per_exp_[instruction] = consumer_add_reduction_elements; } @@ -524,7 +525,8 @@ absl::Status ConstraintPropagator::SeedConstraints( // Output is guaranteed to be non-negative. states_[inst].AddConstraint(ConstraintInterval::Positive()); break; - case HloOpcode::kExp: { + case HloOpcode::kExp: + case HloOpcode::kExpm1: { // Safe domain [-max_log, max_log] prevents floating point overflow. int64_t reduction_elements = GetMaxAddReductionElementsForExp(inst); double max_log = @@ -1599,6 +1601,33 @@ void ConstraintPropagator::PropagateExpApprox( ConstraintInterval{x_min, x_max, /*exclude_zero=*/false}); } +void ConstraintPropagator::PropagateExpm1Approx( + const HloInstruction* instruction, + const ConstraintInterval& output_interval) { + // For Y = expm1(X) = exp(X) - 1 with Y in [y_min, y_max]: + // Since expm1(X) is monotonically strictly increasing on real numbers: + // y_min <= expm1(X) <= y_max <=> ln(y_min + 1) <= X <= ln(y_max + 1). + // + // Since expm1(X) > -1 for all real X: + // If y_max <= -1.0, expm1(X) <= y_max is impossible for real numbers. + if (output_interval.max <= -1.0) { + states_[instruction->operand(0)].AddConstraint( + ConstraintInterval{1.0, -1.0, /*exclude_zero=*/false}); + return; + } + + double x_min = output_interval.min > -1.0 ? std::log1p(output_interval.min) + : ConstraintInterval::kMin; + double x_max = output_interval.max < ConstraintInterval::kMax && + output_interval.max > -1.0 + ? std::log1p(output_interval.max) + : ConstraintInterval::kMax; + bool exclude_zero = output_interval.exclude_zero && + output_interval.min <= 0.0 && output_interval.max >= 0.0; + states_[instruction->operand(0)].AddConstraint( + ConstraintInterval{x_min, x_max, exclude_zero}); +} + void ConstraintPropagator::PropagatePowerApprox( const HloInstruction* instruction, const ConstraintInterval& output_interval) { @@ -1679,6 +1708,9 @@ absl::Status ConstraintPropagator::PropagateConstraintsApprox( case HloOpcode::kExp: PropagateExpApprox(instruction, output_interval); break; + case HloOpcode::kExpm1: + PropagateExpm1Approx(instruction, output_interval); + break; case HloOpcode::kPower: PropagatePowerApprox(instruction, output_interval); break; diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h index 9fa12d6fb24c36..22ff086bd5b2b8 100644 --- a/third_party/xla/xla/tests/constraint_propagator.h +++ b/third_party/xla/xla/tests/constraint_propagator.h @@ -146,6 +146,8 @@ class ConstraintPropagator { const ConstraintInterval& output_interval); void PropagateExpApprox(const HloInstruction* instruction, const ConstraintInterval& output_interval); + void PropagateExpm1Approx(const HloInstruction* instruction, + const ConstraintInterval& output_interval); void PropagatePowerApprox(const HloInstruction* instruction, const ConstraintInterval& output_interval); diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 47d0b515bc739e..b216d3743224c1 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -1503,5 +1503,83 @@ ENTRY main { *module->GetComputationWithName("min_computation")), IdentityElementType::kMaximum); } + +TEST_F(ConstraintPropagatorTest, Expm1SeedConstrainsSafeDomain) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = f32[8,128] parameter(0) + ROOT root = f32[8,128] exponential-minus-one(x) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + // Safe domain [-max_log, max_log] = [-4.0, 4.0] for F32 prevents overflow. + EXPECT_DOUBLE_EQ(x_int.min, -4.0); + EXPECT_DOUBLE_EQ(x_int.max, 4.0); + EXPECT_FALSE(x_int.exclude_zero); +} + +TEST_F(ConstraintPropagatorTest, Expm1BackwardPropagationStrictPositive) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = f32[8,128] parameter(0) + expm1 = f32[8,128] exponential-minus-one(x) + ROOT log = f32[8,128] log(expm1) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + // log(expm1(x)) requires expm1(x) > 0. + // Since expm1(x) > 0 <=> x > log1p(0) = 0, x must be strictly positive. + EXPECT_TRUE(x_int.IsPositiveStrict()); + EXPECT_DOUBLE_EQ(x_int.min, 0.0); +} + +TEST_F(ConstraintPropagatorTest, NanToNumLogFusionConstraints) { + constexpr absl::string_view kHloString = R"hlo( +HloModule fusion.30343, is_scheduled=true + +%fused_computation (param_0: f32[2,6], param_1: f32[2,6]) -> f32[2,6] { + %param_0 = f32[2,6] parameter(0) + %param_1 = f32[2,6] parameter(1) + %expm1 = f32[2,6] exponential-minus-one(%param_0) + %neg = f32[2,6] negate(%expm1) + %add = f32[2,6] add(%param_1, %neg) + %div = f32[2,6] divide(%param_1, %add) + %ne = pred[2,6] compare(%div, %div), direction=NE + %zero = f32[] constant(0) + %broadcast_zero = f32[2,6] broadcast(%zero), dimensions={} + %select_nan = f32[2,6] select(%ne, %broadcast_zero, %div) + ROOT %log = f32[2,6] log(%select_nan) +} + +ENTRY %fusion.30343 (parameter.0: f32[2,6], parameter.1: f32[2,6]) -> f32[2,6] { + %parameter.0 = f32[2,6] parameter(0) + %parameter.1 = f32[2,6] parameter(1) + ROOT %fusion.30343 = f32[2,6] fusion( + %parameter.0, %parameter.1), kind=kLoop, calls=%fused_computation +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto p0_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + auto p1_int = states[module->entry_computation()->parameter_instruction(1)] + .GetConstraintInterval(); + + EXPECT_TRUE(p0_int.IsNegativeStrict()); + EXPECT_DOUBLE_EQ(p0_int.min, -4.0); + EXPECT_TRUE(p1_int.IsPositiveStrict()); +} + } // namespace } // namespace xla From 04ac4ed2b63be6c53756eb77dc1e181bae3fa7ee Mon Sep 17 00:00:00 2001 From: Amit Sabne Date: Fri, 4 Sep 2026 15:39:42 -0700 Subject: [PATCH 04/12] Propagate while-loop layout constraints to conditional subcomputations. This ensures conditional branches inside copy-disabled while loops adopt matching layouts without introducing unnecessary copies. PiperOrigin-RevId: 976527672 --- .../xla/xla/service/layout_assignment.cc | 424 +++++++++++++++++- .../xla/xla/service/layout_assignment.h | 13 + .../xla/xla/service/layout_assignment_test.cc | 78 ++++ 3 files changed, 507 insertions(+), 8 deletions(-) diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index bcc343b0bab574..9cf9e39a46e95e 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -834,6 +834,7 @@ absl::Status LayoutAssignment::AddWhileConstraints( condition_layout.parameter_shape(0))); DCHECK(ShapeUtil::Compatible(body_layout.result_shape(), init->shape())); + bool is_copy_disabled_body = copy_disabled_while_computations_.contains(body); if (body_layout.result_layout() != body_layout.parameter_layout(0)) { VLOG(2) << "Reset %while body parameter layout: body=" << body->name() << " while=" << instruction->name() @@ -843,7 +844,19 @@ absl::Status LayoutAssignment::AddWhileConstraints( body_layout, current_priority_ + kNumberOfPropagationRounds, /*prop_result_layout=*/true, /*prop_parameter_layout=*/true); + if (is_copy_disabled_body) { + while_layout_changed_ = true; + } + } + if (is_copy_disabled_body) { + PropagateWhileLoopLayoutToSubcomputations( + body, body_layout.parameter_layout(0).shape(), + &body_layout.result_layout().shape(), + current_priority_ + kNumberOfPropagationRounds); } + + bool is_copy_disabled_condition = + copy_disabled_while_computations_.contains(condition); if (condition_layout.parameter_layout(0) != body_layout.parameter_layout(0)) { VLOG(2) << "Reset %while condition parameter layout: cond=" << condition->name() << " while=" << instruction->name() @@ -853,12 +866,387 @@ absl::Status LayoutAssignment::AddWhileConstraints( condition_constraint->ResetComputationLayout( condition_layout, current_priority_ + kNumberOfPropagationRounds, /*prop_result_layout=*/true, /*prop_parameter_layout=*/true); + if (is_copy_disabled_condition) { + while_layout_changed_ = true; + } + } + if (is_copy_disabled_condition) { + PropagateWhileLoopLayoutToSubcomputations( + condition, condition_layout.parameter_layout(0).shape(), + &condition_layout.result_layout().shape(), + current_priority_ + kNumberOfPropagationRounds); } ABSL_RETURN_IF_ERROR(SetOperandLayout(body_layout.result_shape(), instruction, 0)); return SetInstructionLayout(body_layout.result_shape(), instruction); } +void LayoutAssignment::PropagateWhileLoopLayoutToSubcomputations( + HloComputation* computation, const Shape& param_shape, + const Shape* result_shape, int64_t priority) { + if (computation == nullptr || computation->num_parameters() == 0) { + return; + } + + auto HasAnyLayout = [](const Shape& shape) -> bool { + bool has = false; + ShapeUtil::ForEachSubshape( + shape, [&has](const Shape& subshape, const ShapeIndex&) { + if (subshape.has_layout()) { + has = true; + } + }); + return has; + }; + + absl::flat_hash_map instruction_layouts; + if (HasAnyLayout(param_shape) && + ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(), + param_shape)) { + instruction_layouts[computation->parameter_instruction(0)] = param_shape; + } + + // Forward propagation from parameter through the computation. + for (HloInstruction* instr : computation->MakeInstructionPostOrder()) { + if (instr == computation->parameter_instruction(0)) { + continue; + } + switch (instr->opcode()) { + case HloOpcode::kGetTupleElement: { + auto it = instruction_layouts.find(instr->operand(0)); + if (it != instruction_layouts.end() && it->second.IsTuple() && + instr->tuple_index() < ShapeUtil::TupleElementCount(it->second)) { + Shape subshape = + ShapeUtil::GetTupleElementShape(it->second, instr->tuple_index()); + if (ShapeUtil::Compatible(instr->shape(), subshape)) { + instruction_layouts[instr] = std::move(subshape); + } + } + break; + } + case HloOpcode::kTuple: { + std::vector subshapes; + subshapes.reserve(instr->operand_count()); + bool any_known = false; + for (int64_t i = 0; i < instr->operand_count(); ++i) { + auto it = instruction_layouts.find(instr->operand(i)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->operand(i)->shape(), it->second)) { + subshapes.push_back(it->second); + any_known = true; + } else { + Shape cleared = instr->operand(i)->shape(); + LayoutUtil::ClearLayout(&cleared); + subshapes.push_back(std::move(cleared)); + } + } + if (any_known) { + Shape tuple_shape = ShapeUtil::MakeTupleShape(subshapes); + if (ShapeUtil::Compatible(instr->shape(), tuple_shape)) { + instruction_layouts[instr] = std::move(tuple_shape); + } + } + break; + } + case HloOpcode::kSelect: { + Shape select_shape; + bool has_select_shape = false; + auto it_true = instruction_layouts.find(instr->operand(1)); + auto it_false = instruction_layouts.find(instr->operand(2)); + if (it_true != instruction_layouts.end() && + HasAnyLayout(it_true->second) && + ShapeUtil::Compatible(instr->shape(), it_true->second)) { + select_shape = it_true->second; + has_select_shape = true; + } else if (it_false != instruction_layouts.end() && + HasAnyLayout(it_false->second) && + ShapeUtil::Compatible(instr->shape(), it_false->second)) { + select_shape = it_false->second; + has_select_shape = true; + } + if (has_select_shape) { + instruction_layouts[instr] = std::move(select_shape); + } + break; + } + case HloOpcode::kCopy: { + auto it = instruction_layouts.find(instr->operand(0)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->shape(), it->second)) { + Shape copy_shape = it->second; + instruction_layouts[instr] = std::move(copy_shape); + } + break; + } + case HloOpcode::kConditional: { + Shape branch_result_shape; + bool has_branch_result_shape = false; + for (int j = 0; j < instr->branch_count(); ++j) { + HloInstruction* branch_arg = instr->mutable_operand(j + 1); + HloComputation* branch_comp = instr->branch_computation(j); + auto it = instruction_layouts.find(branch_arg); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + !ShapeUtil::IsEmptyTuple(it->second) && + ShapeUtil::Compatible(branch_arg->shape(), it->second)) { + PropagateWhileLoopLayoutToSubcomputations( + branch_comp, it->second, /*result_shape=*/nullptr, priority); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr && + branch_constraints->computation_constraint() + .result_layout_is_set()) { + const Shape& res = branch_constraints->computation_layout() + .result_layout() + .shape(); + if (ShapeUtil::Compatible(instr->shape(), res)) { + branch_result_shape = res; + has_branch_result_shape = true; + } + } + } + } + if (!has_branch_result_shape) { + auto it = instruction_layouts.find(instr); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->shape(), it->second)) { + branch_result_shape = it->second; + has_branch_result_shape = true; + } + } + if (has_branch_result_shape) { + instruction_layouts[instr] = branch_result_shape; + for (int j = 0; j < instr->branch_count(); ++j) { + HloComputation* branch_comp = instr->branch_computation(j); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr) { + ComputationLayout branch_layout = + branch_constraints->computation_layout(); + bool prop_branch_result = false; + ShapeUtil::ForEachSubshape( + branch_result_shape, + [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + branch_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_branch_result = true; + } + }); + if (prop_branch_result) { + branch_constraints->mutable_computation_constraint() + ->ResetComputationLayout( + branch_layout, priority, + /*prop_result_layout=*/true, + /*prop_parameter_layout=*/ + branch_constraints->computation_constraint() + .parameter_layout_is_set()); + } + } + } + } + break; + } + case HloOpcode::kWhile: { + HloInstruction* init = instr->mutable_operand(0); + auto it = instruction_layouts.find(init); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(init->shape(), it->second)) { + PropagateWhileLoopLayoutToSubcomputations( + instr->while_body(), it->second, &it->second, priority); + PropagateWhileLoopLayoutToSubcomputations( + instr->while_condition(), it->second, /*result_shape=*/nullptr, + priority); + } + break; + } + case HloOpcode::kCall: { + HloComputation* callee = instr->to_apply(); + if (callee != nullptr && + callee->num_parameters() == instr->operand_count()) { + auto* callee_constraints = mutable_computation_constraints(callee); + if (callee_constraints != nullptr) { + ComputationLayout callee_layout = + callee_constraints->computation_layout(); + bool any_param_set = false; + for (int64_t i = 0; i < instr->operand_count(); ++i) { + auto it = instruction_layouts.find(instr->operand(i)); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(instr->operand(i)->shape(), + it->second)) { + *callee_layout.mutable_parameter_layout(i) = + ShapeLayout(it->second); + any_param_set = true; + } + } + if (any_param_set) { + callee_constraints->mutable_computation_constraint() + ->ResetComputationLayout(callee_layout, priority, + /*prop_result_layout=*/false, + /*prop_parameter_layout=*/true); + } + } + } + break; + } + default: + break; + } + } + + // Backward propagation from result_shape if provided. + if (result_shape != nullptr && HasAnyLayout(*result_shape) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + *result_shape)) { + absl::flat_hash_map root_layouts; + root_layouts[computation->root_instruction()] = *result_shape; + auto post_order = computation->MakeInstructionPostOrder(); + for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) { + HloInstruction* instr = *it; + auto root_it = root_layouts.find(instr); + if (root_it == root_layouts.end() || !HasAnyLayout(root_it->second)) { + continue; + } + Shape current_layout_shape = root_it->second; + switch (instr->opcode()) { + case HloOpcode::kTuple: { + if (current_layout_shape.IsTuple()) { + for (int64_t i = 0; i < instr->operand_count(); ++i) { + if (i < ShapeUtil::TupleElementCount(current_layout_shape)) { + Shape elem_shape = + ShapeUtil::GetTupleElementShape(current_layout_shape, i); + if (ShapeUtil::Compatible(instr->operand(i)->shape(), + elem_shape)) { + root_layouts[instr->operand(i)] = std::move(elem_shape); + } + } + } + } + break; + } + case HloOpcode::kSelect: { + if (ShapeUtil::Compatible(instr->operand(1)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(1)] = current_layout_shape; + } + if (ShapeUtil::Compatible(instr->operand(2)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(2)] = current_layout_shape; + } + break; + } + case HloOpcode::kCopy: { + if (ShapeUtil::Compatible(instr->operand(0)->shape(), + current_layout_shape)) { + root_layouts[instr->operand(0)] = current_layout_shape; + } + break; + } + case HloOpcode::kGetTupleElement: { + HloInstruction* gte_op = instr->mutable_operand(0); + if (gte_op->opcode() == HloOpcode::kConditional && + instr->tuple_index() < + ShapeUtil::TupleElementCount(gte_op->shape())) { + auto cond_it = root_layouts.find(gte_op); + if (cond_it == root_layouts.end()) { + root_layouts[gte_op] = gte_op->shape(); + } + if (root_layouts[gte_op].IsTuple()) { + *ShapeUtil::GetMutableSubshape(&root_layouts[gte_op], + {instr->tuple_index()}) = + current_layout_shape; + } + } + break; + } + case HloOpcode::kConditional: { + if (ShapeUtil::Compatible(instr->shape(), current_layout_shape)) { + for (int j = 0; j < instr->branch_count(); ++j) { + HloComputation* branch_comp = instr->branch_computation(j); + auto* branch_constraints = + mutable_computation_constraints(branch_comp); + if (branch_constraints != nullptr) { + ComputationLayout branch_layout = + branch_constraints->computation_layout(); + bool prop_branch_result = false; + ShapeUtil::ForEachSubshape( + current_layout_shape, + [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + branch_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_branch_result = true; + } + }); + if (prop_branch_result) { + branch_constraints->mutable_computation_constraint() + ->ResetComputationLayout( + branch_layout, priority, + /*prop_result_layout=*/true, + /*prop_parameter_layout=*/ + branch_constraints->computation_constraint() + .parameter_layout_is_set()); + } + } + } + } + break; + } + default: + break; + } + } + } + + // Update this computation's own ComputationLayout. + auto* comp_constraints = mutable_computation_constraints(computation); + if (comp_constraints != nullptr) { + ComputationLayout comp_layout = comp_constraints->computation_layout(); + bool prop_param = false; + if (HasAnyLayout(param_shape) && + ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(), + param_shape)) { + ShapeUtil::ForEachSubshape( + param_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_parameter_layout(0)->ResetLayout( + subshape.layout(), index); + prop_param = true; + } + }); + } + bool prop_result = false; + auto it = instruction_layouts.find(computation->root_instruction()); + if (it != instruction_layouts.end() && HasAnyLayout(it->second) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + it->second)) { + ShapeUtil::ForEachSubshape( + it->second, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_result = true; + } + }); + } else if (result_shape != nullptr && HasAnyLayout(*result_shape) && + ShapeUtil::Compatible(computation->root_instruction()->shape(), + *result_shape)) { + ShapeUtil::ForEachSubshape( + *result_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsArray() && subshape.has_layout()) { + comp_layout.mutable_result_layout()->ResetLayout( + subshape.layout(), index); + prop_result = true; + } + }); + } + if (prop_param || prop_result) { + comp_constraints->mutable_computation_constraint() + ->ResetComputationLayout(comp_layout, priority, prop_result, + prop_param); + } + } +} + absl::Status LayoutAssignment::AddConditionalConstraints( HloInstruction* instruction) { if (computation_layouts_.find(instruction->branch_computation(0)) == @@ -1262,12 +1650,24 @@ absl::Status CheckWhileLayout( auto init_shape = while_inst->operand(0)->shape(); TF_RET_CHECK( condition_computation_layout.parameter_layout(0).MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " cond param layout=" + << condition_computation_layout.parameter_layout(0).ToString() + << " init_shape=" << init_shape.ToString(); TF_RET_CHECK(body_computation_layout.parameter_layout(0).MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " body param layout=" + << body_computation_layout.parameter_layout(0).ToString() + << " init_shape=" << init_shape.ToString(); TF_RET_CHECK(body_computation_layout.result_layout().MatchesLayoutInShape( - init_shape, /*minor_to_major_only=*/true)); - TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape())); + init_shape, /*minor_to_major_only=*/true)) + << " while_inst=" << while_inst->name() << " body result layout=" + << body_computation_layout.result_layout().ToString() + << " init_shape=" << init_shape.ToString(); + TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape())) + << " while_inst=" << while_inst->name() + << " while shape=" << while_inst->shape().ToString() + << " init_shape=" << init_shape.ToString(); return absl::OkStatus(); } @@ -1476,7 +1876,8 @@ absl::Status LayoutAssignment::CopyOperandIfLayoutsDiffer( VLOG(2) << "Operand " << operand->ToString() << " layout does not match " << operand_layout.ToString() << " in " << instruction->ToString(); - if (IsWhileLoopCopyDisabled(*instruction)) { + if (IsWhileLoopCopyDisabled(*instruction) && + !instruction->parent()->caller_instructions(HloOpcode::kWhile).empty()) { HloComputation* comp = instruction->parent(); HloInstruction* param = comp->parameter_instruction(0); ShapeIndex index = {}; @@ -2682,7 +3083,8 @@ absl::Status LayoutAssignment::CalculateComputationLayout( callee->mutable_computation_constraint(); ComputationLayout callee_layout = callee_constraint->computation_layout(); if (callee_constraint->priority() < priority || - conditional_mismatch_.count(callee->computation()) > 0) { + conditional_mismatch_.count(callee->computation()) > 0 || + copy_disabled_while_computations_.contains(callee->computation())) { if (conditional_mismatch_.count(callee->computation()) == 0 && UpdateLayout(result, callee_layout.mutable_result_layout())) { VLOG(2) << "Setting result layout from : " << result->ToString() @@ -2762,7 +3164,8 @@ absl::Status LayoutAssignment::CalculateComputationLayout( } // Reset the layout of the current computation from its body. if (current_priority_ == 0 || - conditional_mismatch_.count(constraints->computation()) > 0) { + conditional_mismatch_.count(constraints->computation()) > 0 || + copy_disabled_while_computations_.contains(constraints->computation())) { ABSL_RETURN_IF_ERROR(SetCalleeLayout( constraints->computation()->root_instruction(), constraints->computation()->parameter_instructions(), constraints, @@ -3240,8 +3643,11 @@ absl::StatusOr LayoutAssignment::RunImpl( // Layouts are allowed to flow naturally in the first round, and any detected // inconsistencies at boundary instructions are resolved with higher-priority // constraints in subsequent rounds. - for (int64_t i = 0; changed || i < kNumberOfPropagationRounds; ++i) { + for (int64_t i = 0; + (changed || i < kNumberOfPropagationRounds) && i < kMaxPropagationRounds; + ++i) { changed = false; + while_layout_changed_ = false; VLOG(1) << "Running " << (i == 0 ? "un" : "") << "constrained pass"; ABSL_RETURN_IF_ERROR(ClearPreviousPassSideEffects(module, execution_threads)); // Layouts are propagated within each computation. In the first round, @@ -3260,6 +3666,7 @@ absl::StatusOr LayoutAssignment::RunImpl( ABSL_ASSIGN_OR_RETURN(bool aliasing_changed, ResolveInputOutputAliasing(module, entry_constraint)); changed |= aliasing_changed; + changed |= while_layout_changed_; } // All logical buffers should have constraints at this point. All that @@ -3450,6 +3857,7 @@ absl::Status LayoutAssignment::Init(HloModule* module) { computation_layouts_.clear(); conditional_mismatch_.clear(); current_priority_ = LayoutConstraint::kBeginningPriority; + while_layout_changed_ = false; // Clear all the copies which have been added, and all the related // instructions (like GTE and tuples). if (!added_copies_.empty()) { diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h index 9c4481fe2a4f4a..daf52cdc9c0c54 100644 --- a/third_party/xla/xla/service/layout_assignment.h +++ b/third_party/xla/xla/service/layout_assignment.h @@ -598,6 +598,7 @@ class LayoutAssignment : public HloModulePass { ComputationLayout& saved_entry_computation_layout() { return saved_entry_computation_layout_; } + virtual bool NegotiateLayout(const HloInstruction* instruction, const Layout& new_layout, const Layout& existing_layout, @@ -803,6 +804,13 @@ class LayoutAssignment : public HloModulePass { absl::Status AddAsyncDoneConstraints(HloInstruction* instruction, LayoutConstraints* constraints); + // Propagates while loop parameter and result layouts to subcomputations (such + // as conditionals) within the while body or condition. + void PropagateWhileLoopLayoutToSubcomputations(HloComputation* computation, + const Shape& param_shape, + const Shape* result_shape, + int64_t priority); + // Propagates layout constraints from the caller instruction into the inner // async sub-computation. // This is the forward propagation step: it takes the layouts of the operands @@ -952,6 +960,7 @@ class LayoutAssignment : public HloModulePass { protected: static constexpr int64_t kNumberOfPropagationRounds = 2; + static constexpr int64_t kMaxPropagationRounds = 6; // Sets up the copy instruction according to the characteristic (sharding, // metadata, ...) of the reference instruction. The index argument is used // when the instruction is a tuple, and in such case the index represents @@ -1072,6 +1081,10 @@ class LayoutAssignment : public HloModulePass { // Stores the set of while computations that have copy disabled. absl::flat_hash_set copy_disabled_while_computations_; + + // Tracks whether while loop parameter/condition layouts changed in the + // current propagation round and require another round to converge. + bool while_layout_changed_ = false; }; } // namespace xla diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc index d0a9091ead93db..659b95b064e256 100644 --- a/third_party/xla/xla/service/layout_assignment_test.cc +++ b/third_party/xla/xla/service/layout_assignment_test.cc @@ -2311,6 +2311,84 @@ ENTRY main { ExpectLayoutIs(ShapeUtil::GetSubshape(body_param_shape, {0}), {0, 1}); } +TEST_F(LayoutAssignmentTest, + RespectsDisableWhileLoopCopiesWithConditionalSubcomputation) { + const char* module_str = R"( +HloModule t + +while_condition { + tuple = (s32[2,8]{1,0}, u32[], pred[]) parameter(0) + i = u32[] get-tuple-element(tuple), index=1 + n = u32[] constant(8) + ROOT predicate = pred[] compare(i, n), direction=LT +} + +branch_true { + arg = s32[2,8]{1,0} parameter(0) + ROOT custom = s32[2,8]{0,1} custom-call(arg), custom_call_target="baz", + operand_layout_constraints={s32[2,8]{0,1}} +} + +branch_false { + ROOT arg = s32[2,8]{1,0} parameter(0) +} + +while_body { + tuple = (s32[2,8]{1,0}, u32[], pred[]) parameter(0) + input = s32[2,8]{1,0} get-tuple-element(tuple), index=0 + i = u32[] get-tuple-element(tuple), index=1 + p = pred[] get-tuple-element(tuple), index=2 + c1 = u32[] constant(1) + i_ = add(i, c1) + cond = s32[2,8]{0,1} conditional(p, input, input), true_computation=branch_true, false_computation=branch_false + ROOT tuple1 = (s32[2,8]{0,1}, u32[], pred[]) tuple(cond, i_, p) +} + +ENTRY main { + input = s32[2,8]{1,0} parameter(0) + c0 = u32[] constant(0) + p0 = pred[] constant(true) + tuple = (s32[2,8]{1,0}, u32[], pred[]) tuple(input, c0, p0) + tuple_ = (s32[2,8]{1,0}, u32[], pred[]) while(tuple), condition=while_condition, body=while_body, frontend_attributes={xla_disable_while_loop_copies="true"} + ROOT output_ = s32[2,8]{1,0} get-tuple-element(tuple_), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); + EXPECT_OK(RunLayoutAssignmentPass(m.get())); + + HloComputation* body = m->GetComputationWithName("while_body"); + ASSERT_NE(body, nullptr); + for (HloInstruction* inst : body->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + + HloComputation* branch_true = m->GetComputationWithName("branch_true"); + ASSERT_NE(branch_true, nullptr); + for (HloInstruction* inst : branch_true->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + ExpectLayoutIs(branch_true->parameter_instruction(0)->shape(), {0, 1}); + ExpectLayoutIs(branch_true->root_instruction()->shape(), {0, 1}); + + HloComputation* branch_false = m->GetComputationWithName("branch_false"); + ASSERT_NE(branch_false, nullptr); + for (HloInstruction* inst : branch_false->instructions()) { + EXPECT_NE(inst->opcode(), HloOpcode::kCopy); + } + ExpectLayoutIs(branch_false->parameter_instruction(0)->shape(), {0, 1}); + ExpectLayoutIs(branch_false->root_instruction()->shape(), {0, 1}); + + HloInstruction* cond = FindInstruction(m.get(), "cond"); + ASSERT_NE(cond, nullptr); + ExpectLayoutIs(cond->shape(), {0, 1}); + + const Shape& body_param_shape = body->parameter_instruction(0)->shape(); + ExpectLayoutIs(ShapeUtil::GetSubshape(body_param_shape, {0}), {0, 1}); + ExpectLayoutIs(ShapeUtil::GetSubshape(body->root_instruction()->shape(), {0}), + {0, 1}); +} + TEST_F(LayoutAssignmentTest, HloBufferLayoutUnconstrained) { const char* module_str = R"( HloModule test From f6957baa1f04902d7f4d0a36e15f9e2403142e7c Mon Sep 17 00:00:00 2001 From: David Dunleavy Date: Fri, 4 Sep 2026 15:58:43 -0700 Subject: [PATCH 05/12] Implement `:riegeli_file_{reader,writer}_factory` with a single .cc file rather than separate `google` and `oss` versions PiperOrigin-RevId: 976535675 --- third_party/xla/xla/service/BUILD | 79 +++---------------- ...ogle.cc => riegeli_file_reader_factory.cc} | 15 +++- .../riegeli_file_reader_factory_oss.cc | 30 ------- ...ogle.cc => riegeli_file_writer_factory.cc} | 15 +++- .../riegeli_file_writer_factory_oss.cc | 30 ------- 5 files changed, 36 insertions(+), 133 deletions(-) rename third_party/xla/xla/service/{riegeli_file_reader_factory_google.cc => riegeli_file_reader_factory.cc} (84%) delete mode 100644 third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc rename third_party/xla/xla/service/{riegeli_file_writer_factory_google.cc => riegeli_file_writer_factory.cc} (84%) delete mode 100644 third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 8ae6293d347548..0b0fd017b209df 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -661,93 +661,34 @@ xla_cc_test( cc_library( name = "riegeli_file_writer_factory", + srcs = ["riegeli_file_writer_factory.cc"], hdrs = ["riegeli_file_writer_factory.h"], + tags = if_google(["ignore_for_dep=third_party/riegeli/bytes/fd_writer.h"]), deps = [ - "@com_google_absl//absl/status", "@com_google_absl//absl/strings:string_view", "@riegeli//riegeli/bytes:writer", + "@tsl//tsl/platform", ] + if_google( - [":riegeli_file_writer_factory_google"], - [":riegeli_file_writer_factory_oss"], + ["@riegeli//riegeli/bytes:file_writer"], + ["@riegeli//riegeli/bytes:fd_writer"], ), ) -cc_library( - name = "riegeli_file_writer_factory_google", - srcs = [ - "riegeli_file_writer_factory.h", - "riegeli_file_writer_factory_google.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - # copybara:uncomment "@riegeli//riegeli/bytes:file_writer", - "@riegeli//riegeli/bytes:writer", - ], - alwayslink = 1, -) - -cc_library( - name = "riegeli_file_writer_factory_oss", - srcs = [ - "riegeli_file_writer_factory.h", - "riegeli_file_writer_factory_oss.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - "@riegeli//riegeli/bytes:fd_writer", - "@riegeli//riegeli/bytes:writer", - ], - alwayslink = 1, -) - cc_library( name = "riegeli_file_reader_factory", + srcs = ["riegeli_file_reader_factory.cc"], hdrs = ["riegeli_file_reader_factory.h"], + tags = if_google(["ignore_for_dep=third_party/riegeli/bytes/fd_reader.h"]), deps = [ "@com_google_absl//absl/strings:string_view", "@riegeli//riegeli/bytes:reader", + "@tsl//tsl/platform", ] + if_google( - [":riegeli_file_reader_factory_google"], - [":riegeli_file_reader_factory_oss"], + ["@riegeli//riegeli/bytes:file_reader"], + ["@riegeli//riegeli/bytes:fd_reader"], ), ) -cc_library( - name = "riegeli_file_reader_factory_google", - srcs = [ - "riegeli_file_reader_factory.h", - "riegeli_file_reader_factory_google.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - # copybara:uncomment "@riegeli//riegeli/bytes:file_reader", - "@riegeli//riegeli/bytes:reader", - ], - alwayslink = 1, -) - -cc_library( - name = "riegeli_file_reader_factory_oss", - srcs = [ - "riegeli_file_reader_factory.h", - "riegeli_file_reader_factory_oss.cc", - ], - tags = ["manual"], - visibility = ["//visibility:private"], - deps = [ - "@com_google_absl//absl/strings:string_view", - "@riegeli//riegeli/bytes:fd_reader", - "@riegeli//riegeli/bytes:reader", - ], - alwayslink = 1, -) - xla_cc_test( name = "dump_test", srcs = ["dump_test.cc"], diff --git a/third_party/xla/xla/service/riegeli_file_reader_factory_google.cc b/third_party/xla/xla/service/riegeli_file_reader_factory.cc similarity index 84% rename from third_party/xla/xla/service/riegeli_file_reader_factory_google.cc rename to third_party/xla/xla/service/riegeli_file_reader_factory.cc index 1f60b9689307d5..02f2f9910afa6e 100644 --- a/third_party/xla/xla/service/riegeli_file_reader_factory_google.cc +++ b/third_party/xla/xla/service/riegeli_file_reader_factory.cc @@ -13,18 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/service/riegeli_file_reader_factory.h" + #include #include "absl/strings/string_view.h" -#include "riegeli/bytes/file_reader.h" #include "riegeli/bytes/reader.h" -#include "xla/service/riegeli_file_reader_factory.h" +#include "tsl/platform/platform.h" + +#if TSL_IS_IN_OSS +#include "riegeli/bytes/fd_reader.h" +#else +#include "riegeli/bytes/file_reader.h" +#endif namespace xla { std::unique_ptr CreateRiegeliFileReader( absl::string_view filename) { +#if TSL_IS_IN_OSS + return std::make_unique>(filename); +#else return std::make_unique>(filename); +#endif } } // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc b/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc deleted file mode 100644 index 479bc0939f2473..00000000000000 --- a/third_party/xla/xla/service/riegeli_file_reader_factory_oss.cc +++ /dev/null @@ -1,30 +0,0 @@ -/* Copyright 2026 The OpenXLA 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. -==============================================================================*/ - -#include - -#include "absl/strings/string_view.h" -#include "riegeli/bytes/fd_reader.h" -#include "riegeli/bytes/reader.h" -#include "xla/service/riegeli_file_reader_factory.h" - -namespace xla { - -std::unique_ptr CreateRiegeliFileReader( - absl::string_view filename) { - return std::make_unique>(filename); -} - -} // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_writer_factory_google.cc b/third_party/xla/xla/service/riegeli_file_writer_factory.cc similarity index 84% rename from third_party/xla/xla/service/riegeli_file_writer_factory_google.cc rename to third_party/xla/xla/service/riegeli_file_writer_factory.cc index c236236b66d140..4b72cf47b0507d 100644 --- a/third_party/xla/xla/service/riegeli_file_writer_factory_google.cc +++ b/third_party/xla/xla/service/riegeli_file_writer_factory.cc @@ -13,18 +13,29 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/service/riegeli_file_writer_factory.h" + #include #include "absl/strings/string_view.h" -#include "riegeli/bytes/file_writer.h" #include "riegeli/bytes/writer.h" -#include "xla/service/riegeli_file_writer_factory.h" +#include "tsl/platform/platform.h" + +#if TSL_IS_IN_OSS +#include "riegeli/bytes/fd_writer.h" +#else +#include "riegeli/bytes/file_writer.h" +#endif namespace xla { std::unique_ptr CreateRiegeliFileWriter( absl::string_view filename) { +#if TSL_IS_IN_OSS + return std::make_unique>(filename); +#else return std::make_unique>(filename); +#endif } } // namespace xla diff --git a/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc b/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc deleted file mode 100644 index 19fb7e9297f897..00000000000000 --- a/third_party/xla/xla/service/riegeli_file_writer_factory_oss.cc +++ /dev/null @@ -1,30 +0,0 @@ -/* Copyright 2026 The OpenXLA 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. -==============================================================================*/ - -#include - -#include "absl/strings/string_view.h" -#include "riegeli/bytes/fd_writer.h" -#include "riegeli/bytes/writer.h" -#include "xla/service/riegeli_file_writer_factory.h" - -namespace xla { - -std::unique_ptr CreateRiegeliFileWriter( - absl::string_view filename) { - return std::make_unique>(filename); -} - -} // namespace xla From dfd475104196efb5bd95ac4d730437e6e9a374f6 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Fri, 4 Sep 2026 16:19:14 -0700 Subject: [PATCH 06/12] Refactor away LocalDeviceState* usage on PjRtStreamExecutorDevice. This should always be available from RawClient with local device id. PiperOrigin-RevId: 976544198 --- .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc | 78 ++++++------ .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h | 4 +- .../gpu/se_gpu_pjrt_client_multi_gpu_test.cc | 10 +- .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc | 25 ++-- .../pjrt/se/pjrt_stream_executor_client.cc | 119 ++++++++---------- .../xla/pjrt/se/pjrt_stream_executor_client.h | 56 +++++---- .../se/pjrt_stream_executor_client_test.cc | 15 ++- 7 files changed, 146 insertions(+), 161 deletions(-) diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc index ca8451d79ad792..95028d7f264cda 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc @@ -235,10 +235,10 @@ CreateSEGpuTopology(absl::string_view platform_name, } static se::StreamExecutor* GetFirstExecutor( - const std::vector>& devices) { - for (const auto& d : devices) { - if (auto* local_device_state = d.get()->local_device_state()) { - return local_device_state->executor(); + absl::Span> local_device_states) { + for (const auto& state : local_device_states) { + if (state != nullptr) { + return state->executor(); } } return nullptr; @@ -249,15 +249,16 @@ namespace { // Derives the platform version string (e.g. "cuda 12080") from the runtime // device description, independent of compile-time macros. static std::string GpuPlatformVersionFromDevices( + const StreamExecutorGpuRawClient* raw_client, absl::Span> devices) { for (const std::unique_ptr& device : devices) { - auto* se_device = device.get(); - LocalDeviceState* local_device_state = se_device->local_device_state(); - if (local_device_state == nullptr) { + auto local_device_state = + raw_client->GetLocalDeviceState(device->local_device_id()); + if (!local_device_state.ok()) { continue; } const se::DeviceDescription& desc = - local_device_state->executor()->GetDeviceDescription(); + local_device_state.value()->executor()->GetDeviceDescription(); const se::SemanticVersion v = desc.runtime_version(); const se::GpuComputeCapability& cc = desc.gpu_compute_capability(); if (cc.rocm_compute_capability() != nullptr) { @@ -306,13 +307,6 @@ void StreamExecutorGpuRawClient::UpdateCompileOptionsTopology( // helpers. namespace { -// Get the local device state for a given PjRtDevice. -absl::StatusOr GetLocalDeviceState(PjRtDevice* device) { - PjRtStreamExecutorDevice* pjrt_se_device = - absl::down_cast(device); - return pjrt_se_device->GetLocalDeviceState(); -} - // Creates a communicator for a cross-host transfer; used by the original // cross-host transfers API. absl::StatusOr> CreateTransferCommunicator( @@ -627,31 +621,29 @@ StreamExecutorGpuRawClient::CrossHostTransferBuffers( // Get the local_device_state and use it to schedule transfers. Fail // transfers early if we cannot get the local_device_state. - absl::StatusOr local_device_state = - absl::down_cast(device) - ->GetLocalDeviceState(); - if (!local_device_state.ok()) { - SetEventAsError(transfer_event, local_device_state.status()); + auto local_device_state_or = GetLocalDeviceState(device->local_device_id()); + if (!local_device_state_or.ok()) { + SetEventAsError(transfer_event, local_device_state_or.status()); continue; } + LocalDeviceState* local_device_state = *local_device_state_or; // Launch ScheduleTransfersOnLocalDevice on either the async dispatch thread // of the calling thread. - if ((*local_device_state)->async_dispatch_thread()) { - (*local_device_state) - ->async_dispatch_thread() - ->Schedule(tsl::WithCurrentContext( + if (local_device_state->async_dispatch_thread()) { + local_device_state->async_dispatch_thread()->Schedule( + tsl::WithCurrentContext( [this, local_device_state, device_id, transfer_dependencies, curr_transfer_specs = std::move(curr_transfer_specs), transfer_event = std::move(transfer_event)]() mutable { - ScheduleTransfersOnLocalDevice(*local_device_state, device_id, + ScheduleTransfersOnLocalDevice(local_device_state, device_id, std::move(transfer_event), std::move(transfer_dependencies), std::move(curr_transfer_specs)); })); } else { ScheduleTransfersOnLocalDevice( - *local_device_state, device_id, std::move(transfer_event), + local_device_state, device_id, std::move(transfer_event), transfer_dependencies, std::move(curr_transfer_specs)); } } @@ -1748,8 +1740,13 @@ absl::StatusOr BuildDistributedDevices( local_device_states_vec.push_back(std::move(it->second)); } + if (local_device != nullptr) { + CHECK_EQ(local_device->local_device_id().value(), + device_proto.local_device_ordinal()); + } auto device = std::make_unique( - device_proto.global_device_id(), local_device, device_proto.name(), + device_proto.global_device_id(), + /*is_addressable=*/local_device != nullptr, device_proto.name(), device_proto.vendor(), device_proto.compute_capability(), device_proto.core_count(), device_proto.device_memory_bytes_limit(), device_proto.shared_memory_per_block_optin(), @@ -1793,12 +1790,12 @@ absl::StatusOr BuildDistributedDevices( } StreamExecutorGpuDevice::StreamExecutorGpuDevice( - int id, LocalDeviceState* local_device_state, std::string device_kind, + int id, bool is_addressable, std::string device_kind, std::string device_vendor, std::string compute_capability, int core_count, int64_t device_memory_bytes_limit, int64_t shared_memory_per_block_optin, int local_device_id, int process_index, int process_index_in_partition, int partition_index, int numa_node, std::string fabric_uuid) - : PjRtStreamExecutorDevice(id, local_device_state, local_device_id, + : PjRtStreamExecutorDevice(id, is_addressable, local_device_id, process_index, process_index_in_partition, partition_index, std::move(device_kind)) { VLOG(1) << absl::StreamFormat( @@ -1894,7 +1891,8 @@ std::unique_ptr MakeStreamExecutorGpuClient( std::shared_ptr kv_store, std::shared_ptr topology, std::optional num_processes) { - std::string platform_version = GpuPlatformVersionFromDevices(devices); + std::string platform_version = + GpuPlatformVersionFromDevices(raw_client.get(), devices); PjRtPluginAttributes attrs; attrs.pjrt_c_api_major_version = 0; attrs.pjrt_c_api_minor_version = 0; @@ -2096,16 +2094,16 @@ absl::StatusOr> GetStreamExecutorGpuClient( ABSL_ASSIGN_OR_RETURN(std::shared_ptr gpu_topology, GpuTopology::FromProto(devices_and_topology.topology)); - auto se_gpu_topology = - CreateSEGpuTopology(pjrt_platform_name, std::move(gpu_topology), - GetFirstExecutor(devices_and_topology.devices)); + se::StreamExecutor* first_executor = + GetFirstExecutor(devices_and_topology.local_device_states); + auto se_gpu_topology = CreateSEGpuTopology( + pjrt_platform_name, std::move(gpu_topology), first_executor); auto raw_client = std::make_unique( tsl::Fingerprint64(pjrt_platform_name), std::move(devices_and_topology.local_device_states), std::move(allocator), xla_client, std::move(host_memory_allocator), options.should_stage_host_to_device_transfers, - /*async_work_runner=*/nullptr, - GetFirstExecutor(devices_and_topology.devices), kv_store, + /*async_work_runner=*/nullptr, first_executor, kv_store, preallocate_device_memory, options.abort_collectives_on_failure, std::move(gpu_run_options), std::move(memory_registration)); VLOG(1) << absl::StreamFormat( @@ -2156,16 +2154,16 @@ absl::StatusOr> GetSharedStreamExecutorGpuClient( ABSL_ASSIGN_OR_RETURN(auto gpu_topology, absl::StatusOr>( GpuTopology::FromProto(devices_and_topology.topology))); - auto se_gpu_topology = - CreateSEGpuTopology(platform_name, std::move(gpu_topology), - GetFirstExecutor(devices_and_topology.devices)); + se::StreamExecutor* first_executor = + GetFirstExecutor(devices_and_topology.local_device_states); + auto se_gpu_topology = CreateSEGpuTopology( + platform_name, std::move(gpu_topology), first_executor); auto raw_client = std::make_unique( tsl::Fingerprint64(platform_name), std::move(devices_and_topology.local_device_states), std::move(allocator), local_client, std::move(host_memory_allocator), /*should_stage_host_to_device_transfers=*/true, - /*async_work_runner=*/nullptr, - GetFirstExecutor(devices_and_topology.devices), kv_store, + /*async_work_runner=*/nullptr, first_executor, kv_store, /*cache_fabric_handles=*/false, /*abort_collectives_on_failure=*/false, std::move(gpu_run_options)); VLOG(1) << absl::StreamFormat( diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h index e79b15e8b6ed91..3467c593887e28 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h @@ -76,8 +76,8 @@ namespace xla { class StreamExecutorGpuDevice : public PjRtStreamExecutorDevice { public: - StreamExecutorGpuDevice(int id, LocalDeviceState* local_device_state, - std::string device_kind, std::string device_vendor, + StreamExecutorGpuDevice(int id, bool is_addressable, std::string device_kind, + std::string device_vendor, std::string compute_capability, int core_count, int64_t device_memory_bytes_limit, int64_t shared_memory_per_block_optin, diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc index b3c4969501fa5f..e31a33843a4801 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc @@ -1678,11 +1678,11 @@ absl::Status InterProcessCollectiveInitTestBody(int rank_id) { // executor's collective memory allocator into the selected collectives // backend (e.g. MORI ShmemMalloc). With inert backend stubs the allocation // may return null; we only log the outcome and do not fail the test. - auto* se_device = absl::down_cast( - client->addressable_devices()[0]); - TF_RET_CHECK(se_device != nullptr); - LocalDeviceState* local_device_state = se_device->local_device_state(); - TF_RET_CHECK(local_device_state != nullptr); + auto* se_client = absl::down_cast(client.get()); + TF_RET_CHECK(se_client != nullptr); + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device_state, + se_client->raw_client()->GetLocalDeviceState( + client->addressable_devices()[0]->local_device_id())); se::StreamExecutor* executor = local_device_state->executor(); constexpr uint64_t kCollectiveBytes = 1024; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index 593f3cb71a3733..86e4bd4d8095e9 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -229,13 +229,13 @@ TEST(StreamExecutorGpuClientTest, AsyncResultDefinitionEventUsesAsyncStream) { ASSERT_OK_AND_ASSIGN(auto definition, GetDefinitionStreamInfo(results[0][0].get())); - auto* se_device = absl::down_cast( - client->addressable_devices().front()); - intptr_t compute_stream = - reinterpret_cast(se_device->local_device_state() - ->compute_stream() - ->platform_specific_handle() - .stream); + auto* se_client = absl::down_cast(client.get()); + TF_ASSERT_OK_AND_ASSIGN( + LocalDeviceState * local_device_state, + se_client->raw_client()->GetLocalDeviceState( + client->addressable_devices().front()->local_device_id())); + intptr_t compute_stream = reinterpret_cast( + local_device_state->compute_stream()->platform_specific_handle().stream); EXPECT_NE(definition.stream, compute_stream); ASSERT_OK_AND_ASSIGN(auto literal, results[0][0]->ToLiteral().Await()); @@ -2568,12 +2568,11 @@ TEST(StreamExecutorGpuClientTest, EventCaching) { absl::down_cast(client.get()) ->async_work_runner(); const auto& device = client->addressable_devices()[0]; - // TODO(b/b/482307468) Switch to absl::down_cast after upgrade. - [[deprecated( - "remove after absl upgrade")]] LocalDeviceState* local_device_state = - absl::down_cast(device) - ->local_device_state(); - ASSERT_TRUE(local_device_state != nullptr); + TF_ASSERT_OK_AND_ASSIGN( + LocalDeviceState * local_device_state, + absl::down_cast(client.get()) + ->raw_client() + ->GetLocalDeviceState(device->local_device_id())); size_t sync_point0 = local_device_state->GetNextComputeStreamSyncPoint(); TF_ASSERT_OK_AND_ASSIGN(auto event0, local_device_state->GetEventForComputeStreamSyncPoint( diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index e680c677ef0429..f433928799737b 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -235,14 +235,6 @@ absl::string_view PjRtStreamExecutorDevice::platform_name() const { return client_->platform_name(); } -absl::StatusOr -PjRtStreamExecutorDevice::GetLocalDeviceState() const { - if (local_device_state_ != nullptr) { - return local_device_state_; - } - return InvalidArgument("Device %s is not a local device.", DebugString()); -} - absl::StatusOr DevicesToDeviceAssignment( absl::Span> devices) { if (devices.empty()) { @@ -427,10 +419,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( bool retry_on_oom, tsl::AsyncValueRef allocate_after) { CHECK(allocate_after == nullptr) << "allocate_after is not supported for PjRtStreamExecutorClient."; - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); PjRtMemorySpace* default_memory_space = device->default_memory_space().value_or(nullptr); auto layout_memory_space = Layout::kDefaultMemorySpace; @@ -477,10 +468,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBufferForExecute( PjRtMemorySpace* memory_space, size_t on_device_bytes_count, bool retry_on_oom) { - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto mem = RawSEDeviceMemory::CreateDelayedMemory(); return tsl::MakeRef( this, memory_space, local_device, std::move(mem), on_device_bytes_count); @@ -491,10 +481,9 @@ absl::StatusOr( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto raw_buffer = tsl::MakeRef( this, memory_space, local_device, tsl::AsyncValueRef(buffer_promise), @@ -614,10 +603,9 @@ PjRtStreamExecutorClient::AllocateLinearizeDest( absl::StatusOr> PjRtStreamExecutorRawClient::CreateLinkedEventPromise( PjRtMemorySpace* memory_space, absl::string_view debug_info) { - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto result = tsl::MakeRef( this, local_device, async_work_runner()); PjRtDeviceEventRef event = result->event().CopyRef(); @@ -637,10 +625,9 @@ absl::StatusOr PjRtStreamExecutorRawClient::CreateDeviceEvent(PjRtMemorySpace* memory_space, Future<> dependency) { auto definition_event = BufferSequencingEvent::Create(async_work_runner()); - auto* device = tensorflow::down_cast( - memory_space->devices()[0]); + PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - device->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); dependency.OnReady([definition_event = definition_event.CopyRef(), local_device, this](absl::Status status) mutable { if (!status.ok()) { @@ -704,8 +691,7 @@ PjRtStreamExecutorRawClient::ImportForeignMemory( CHECK_EQ(memory_space->devices().size(), 1); auto* device = memory_space->devices().front(); ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - tensorflow::down_cast(device) - ->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto buffer = RawSEDeviceMemory::CreateForeign( se::DeviceAddressBase(device_ptr, size), std::move(on_delete_callback)); @@ -720,8 +706,7 @@ PjRtStreamExecutorRawClient::CreateDeviceEventForStream( CHECK_EQ(memory_space->devices().size(), 1); auto* device = memory_space->devices().front(); ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - absl::down_cast(device) - ->GetLocalDeviceState()); + GetLocalDeviceState(device->local_device_id())); auto definition_event = BufferSequencingEvent::Create(this->async_work_runner()); @@ -741,10 +726,7 @@ PjRtStreamExecutorExecutableLoadState::LoadRawExecutable( const ExecuteOptions& options, size_t host_callback_idx, xla::RunId run_id, DeviceAndAssignment device_and_assign, int attempt) { PjRtDevice* device = device_and_assign.device; - int device_ordinal = absl::down_cast(device) - ->local_device_state() - ->local_device_id() - .value(); + int device_ordinal = device->local_device_id().value(); auto se_executable = std::move(executable).Cast(); const CompileOptions& compile_options = se_executable->compile_options(); ABSL_ASSIGN_OR_RETURN(auto local_exec, @@ -762,30 +744,40 @@ PjRtStreamExecutorExecutableLoadState::LoadRawExecutable( std::move(se_executable), raw_client()); } +absl::Status PjRtStreamExecutorRawClient::TransferToInfeed( + LocalDeviceId local_device_id, const LiteralSlice& literal) { + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + GetLocalDeviceState(local_device_id)); + return NeverRunOnFiber(async_work_runner(), [&]() { + return local_device->client()->TransferToInfeedLocal( + literal, local_device->local_hardware_id().value()); + }); +} + +absl::Status PjRtStreamExecutorRawClient::TransferFromOutfeed( + LocalDeviceId local_device_id, MutableBorrowingLiteral literal) { + VLOG(1) << "PjRtStreamExecutorRawClient::TransferFromOutfeed"; + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + GetLocalDeviceState(local_device_id)); + return NeverRunOnFiber(async_work_runner(), [&]() { + return local_device->client()->TransferFromOutfeedLocal( + local_device->local_hardware_id().value(), literal); + }); +} + // Transfer the given literal to the infeed queue of the given local device. absl::Status PjRtStreamExecutorDevice::TransferToInfeed( const LiteralSlice& literal) { - // Only support infeed to local device. - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); - return NeverRunOnFiber( - tensorflow::down_cast(client_) - ->async_work_runner(), - [&]() { - return local_device->client()->TransferToInfeedLocal( - literal, local_device->local_hardware_id().value()); - }); + return absl::down_cast(client_) + ->raw_client() + ->TransferToInfeed(local_device_id(), literal); } absl::Status PjRtStreamExecutorDevice::TransferFromOutfeed( MutableBorrowingLiteral literal) { - VLOG(1) << "PjRtStreamExecutorDevice::TransferFromOutfeed"; - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); - return NeverRunOnFiber( - absl::down_cast(client_)->async_work_runner(), - [&]() { - return local_device->client()->TransferFromOutfeedLocal( - local_device->local_hardware_id().value(), literal); - }); + return absl::down_cast(client_) + ->raw_client() + ->TransferFromOutfeed(local_device_id(), literal); } void PjRtStreamExecutorDevice::AttachMemorySpace(PjRtMemorySpace* memory_space, @@ -846,7 +838,10 @@ PjRtStreamExecutorDevice::memory_space_by_kind_id(int id) const { absl::StatusOr PjRtStreamExecutorDevice::GetStreamForExternalReadyEvents() const { - ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState()); + ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, + absl::down_cast(client_) + ->raw_client() + ->GetLocalDeviceState(local_device_id())); se::Stream* stream = local_device->GetExternalReadyEventStream(); void* raw_stream = stream->platform_specific_handle().stream; if (raw_stream == nullptr) { @@ -1407,13 +1402,9 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( PjRtDeviceEventRefVector extra_deps, PjRtDeviceEventRefVector control_deps, bool is_predetermined_error, bool fill_future) && { const uint64_t start_time_usecs = tsl::Env::Default()->NowMicros(); - int device_ordinal = tensorflow::down_cast(device_) - ->local_device_state() - ->local_device_id() - .value(); + int device_ordinal = device_->local_device_id().value(); LocalDeviceState* device_state = - tensorflow::down_cast(device_) - ->local_device_state(); + raw_client_->device_state(device_->local_device_id()); const CompileOptions& compile_options = se_executable_->compile_options(); tsl::profiler::TraceMe trace([&] { @@ -1693,9 +1684,6 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( }; auto definition_event = [&]() -> PjRtDeviceEventRef { - LocalDeviceState* device_state = - tensorflow::down_cast(device) - ->local_device_state(); se::Stream* stream = device_state->compute_stream(); if (!result_buffer_or_status.ok()) { @@ -1805,8 +1793,7 @@ PjRtStreamExecutorRawLoadedExecutable::Execute( void PjRtStreamExecutorClient::LaunchOnDevice( PjRtDevice* device, absl::AnyInvocable execute_fn) const { const LocalDeviceState& device_state = - *tensorflow::down_cast(device) - ->local_device_state(); + this->device_state(device->local_device_id().value()); device_state.execute_thread()->Schedule( tsl::WithCurrentContext(std::move(execute_fn))); } @@ -2190,8 +2177,7 @@ PjRtStreamExecutorClient::Load(std::shared_ptr executable, LoadInternal(std::move(executable), load_options, /*dump=*/false); for (const PjRtDevice* device : addressable_devices()) { LocalDeviceState* local_device_state = - tensorflow::down_cast(device) - ->local_device_state(); + raw_client()->device_state(device->local_device_id()); raw_client()->RecordMemoryStats(local_device_state); } return loaded_executable; @@ -2202,13 +2188,12 @@ bool PjRtStreamExecutorClient::IsHostMemoryPinned(const void* ptr, if (addressable_devices().empty()) { return false; } - auto* device = tensorflow::down_cast( - addressable_devices()[0]); - auto status_or_device_state = device->GetLocalDeviceState(); - if (!status_or_device_state.ok()) { + auto local_device_state = raw_client()->GetLocalDeviceState( + addressable_devices()[0]->local_device_id()); + if (!local_device_state.ok()) { return false; } - return status_or_device_state.value() + return local_device_state.value() ->compute_stream() ->parent() ->IsHostMemoryPinned(ptr, size); diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h index 359820015aedd9..81a3823a8f0adf 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h @@ -106,22 +106,20 @@ class StreamExecutorExecutable; class PjRtStreamExecutorDevice : public PjRtDevice { public: - PjRtStreamExecutorDevice(int id, LocalDeviceState* local_device_state, - int local_device_id, int process_index, - int process_index_in_partition, int partition_index, - std::string device_kind) - : local_device_id_(local_device_id), - local_hardware_id_(local_device_state - ? local_device_state->local_hardware_id() - : LocalChipId(-1)), - local_device_state_(local_device_state), + PjRtStreamExecutorDevice(int id, bool is_addressable, int local_device_id, + int process_index, int process_index_in_partition, + int partition_index, std::string device_kind, + LocalChipId local_hardware_id = LocalChipId(-1)) + : is_addressable_(is_addressable), + local_device_id_(local_device_id), + local_hardware_id_(local_hardware_id.value() != -1 + ? local_hardware_id + : (is_addressable ? LocalChipId(local_device_id) + : LocalChipId(-1))), description_(id, local_device_id_.value(), process_index, process_index_in_partition, partition_index, - std::move(device_kind)) { - if (local_device_state_ != nullptr) { - CHECK_EQ(local_device_state_->local_device_id(), local_device_id_); - } - } + std::move(device_kind)) {} + ~PjRtStreamExecutorDevice() override = default; // Must set client exactly once. @@ -155,7 +153,7 @@ class PjRtStreamExecutorDevice : public PjRtDevice { PjRtClient* client() const override { return client_; } - bool IsAddressable() const override { return local_device_state_ != nullptr; } + bool IsAddressable() const override { return is_addressable_; } LocalDeviceId local_device_id() const override { return local_device_id_; } @@ -166,16 +164,6 @@ class PjRtStreamExecutorDevice : public PjRtDevice { return attributes_; } - // If this is a device local to this host, returns a LocalDeviceState object - // that can be used to manipulate the device. Returns nullptr if the device is - // not local to this host. - LocalDeviceState* local_device_state() const { return local_device_state_; } - - // If this is a device local to this host, returns a LocalDeviceState object - // that can be used to manipulate the device. Returns an error if the device - // is not local to this host. - absl::StatusOr GetLocalDeviceState() const; - absl::Status TransferToInfeed(const LiteralSlice& literal) override; absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) override; @@ -201,9 +189,9 @@ class PjRtStreamExecutorDevice : public PjRtDevice { } private: + const bool is_addressable_; const LocalDeviceId local_device_id_; const LocalChipId local_hardware_id_; - LocalDeviceState* local_device_state_ = nullptr; PjRtStreamExecutorDeviceDescription description_; absl::flat_hash_map attributes_; PjRtClient* client_ = nullptr; @@ -263,6 +251,16 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { return it != local_device_states_by_id_.end() ? it->second : nullptr; } + absl::StatusOr GetLocalDeviceState( + LocalDeviceId local_device_id) const { + LocalDeviceState* state = device_state(local_device_id); + if (state == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "Device ", local_device_id.value(), " is not a local device.")); + } + return state; + } + gpu::GpuExecutableRunOptions* gpu_run_options() const { return gpu_run_options_.get(); } @@ -362,6 +360,12 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { absl::StatusOr CreateDeviceEventForStream( PjRtMemorySpace* memory_space, std::intptr_t stream) override; + absl::Status TransferToInfeed(LocalDeviceId local_device_id, + const LiteralSlice& literal) override; + + absl::Status TransferFromOutfeed(LocalDeviceId local_device_id, + MutableBorrowingLiteral literal) override; + virtual void UpdateCompileOptionsTopology( const PjRtTopologyDescription& topology, CompileOptions* options) const {} diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index 9d08e8d509e83c..baf79d2c5a3fe5 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -169,10 +169,9 @@ MakeTestPjRtStreamExecutorClient( std::unique_ptr gpu_run_options = nullptr, std::shared_ptr kv_store = nullptr) { se::StreamExecutor* first_executor = nullptr; - for (const auto& dev : devices) { - if (dev->IsAddressable() && dev->local_device_state() != nullptr && - dev->local_device_state()->compute_stream() != nullptr) { - first_executor = dev->local_device_state()->compute_stream()->parent(); + for (const auto& state : local_device_states) { + if (state != nullptr && state->compute_stream() != nullptr) { + first_executor = state->compute_stream()->parent(); break; } } @@ -209,7 +208,7 @@ absl::StatusOr> GetClient() { int local_device_id = local_device_states.back()->local_device_id().value(); std::vector> devices; devices.emplace_back(std::make_unique( - 0, local_device_states.back().get(), local_device_id, /*process_index=*/0, + 0, /*is_addressable=*/true, local_device_id, /*process_index=*/0, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); std::vector> memory_spaces; memory_spaces.emplace_back(std::make_unique( @@ -261,7 +260,7 @@ absl::StatusOr> GetClientWithDevices( /*allow_event_reuse=*/false, /*use_callback_stream=*/false)); int local_device_id = local_device_states.back()->local_device_id().value(); devices.emplace_back(std::make_unique( - i, local_device_states.back().get(), local_device_id, + i, /*is_addressable=*/true, local_device_id, /*process_index=*/0, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); memory_spaces.emplace_back(std::make_unique( @@ -271,7 +270,7 @@ absl::StatusOr> GetClientWithDevices( } for (int i = num_addressable_devices; i < total_devices; ++i) { devices.emplace_back(std::make_unique( - i, /*local_device_state=*/nullptr, /*local_device_id=*/-1, + i, /*is_addressable=*/false, /*local_device_id=*/-1, /*process_index=*/1, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); } @@ -485,7 +484,7 @@ TEST(PjRtStreamExecutorClientTest, ExecutePortableRemoteDevice) { *client, shape, [](XlaBuilder& builder) {}, compile_options)); auto remote_device = std::make_unique( - 1, /*local_device_state=*/nullptr, /*local_device_id=*/-1, + 1, /*is_addressable=*/false, /*local_device_id=*/-1, /*process_index=*/1, /*process_index_in_partition=*/1, /*partition_index=*/0, "cpu"); remote_device->SetClient(client.get()); From fe8174a37a767367a7de85b741edfac1185a2834 Mon Sep 17 00:00:00 2001 From: David Pizzuto Date: Fri, 4 Sep 2026 16:34:24 -0700 Subject: [PATCH 07/12] pjrt_client: Add IsCApi method. Today the way to infer whether the C API is used is by checking platform_version strings, which is indirect and also makes it difficult to match versions between C API and C++ API objects. Notably, topology serializations of otherwise identical protos would mismatch because of this. PJRT_API_MINOR is not incremented because this method does not actually go over the C API. PiperOrigin-RevId: 976549887 --- third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc | 2 ++ third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h | 2 ++ .../xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc | 7 +++++++ third_party/xla/xla/pjrt/pjrt_client.cc | 2 ++ third_party/xla/xla/pjrt/pjrt_client.h | 3 +++ third_party/xla/xla/pjrt/pjrt_client_test.cc | 5 +++++ 6 files changed, 21 insertions(+) diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc index 7ea1ed00f02ba1..0c2fed894134ec 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc @@ -188,6 +188,8 @@ PjRtCApiClient::PjRtCApiClient( LOG(INFO) << "PjRtCApiClient created."; } +bool PjRtCApiClient::IsCApi() const { return true; } + void PjRtCApiClient::InitDevicesAndMemorySpaces() { // Initialize devices. PJRT_Client_Devices_Args devices_args; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h index 4b4da58f2ffdcc..c60c8df0da1613 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h @@ -372,6 +372,8 @@ class PjRtCApiClient : public PjRtClient { const PJRT_Api* c_api, PJRT_Client* c_client, std::unique_ptr<::pjrt::PJRT_KeyValueCallbackData> kv_callback_data); + bool IsCApi() const override; + int process_index() const override; int device_count() const override; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc index b7efe0487da9d0..a9c226fde55181 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc @@ -1137,5 +1137,12 @@ TEST(PjRtCApiClientTest, MakeCanonicalShapeForMemorySpace) { specific_layout.minor_to_major()); } +TEST(PjRtCApiClientTest, IsCApi) { + SetUpCpuPjRtApi(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr client, + GetCApiClient("cpu")); + EXPECT_TRUE(client->IsCApi()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/pjrt/pjrt_client.cc b/third_party/xla/xla/pjrt/pjrt_client.cc index efca03ade466aa..28c34b4ca7ef62 100644 --- a/third_party/xla/xla/pjrt/pjrt_client.cc +++ b/third_party/xla/xla/pjrt/pjrt_client.cc @@ -119,6 +119,8 @@ PjRtMemorySpaceCApiDelegator::PjRtMemorySpaceCApiDelegator( PjRtBuffer::ExternalReference::~ExternalReference() = default; +bool PjRtClient::IsCApi() const { return false; } + absl::StatusOr PjRtClient::UnsafeBufferPointer( PjRtBuffer* buffer) { if (buffer->on_device_shape().IsTuple()) { diff --git a/third_party/xla/xla/pjrt/pjrt_client.h b/third_party/xla/xla/pjrt/pjrt_client.h index 27ed24b15259e3..8ad889abe2dabc 100644 --- a/third_party/xla/xla/pjrt/pjrt_client.h +++ b/third_party/xla/xla/pjrt/pjrt_client.h @@ -558,6 +558,9 @@ class PjRtClient { virtual ~PjRtClient() = default; + // Whether this client uses the C API. + virtual bool IsCApi() const; + // Return the process index of this client. Always 0 in single-process // settings. virtual int process_index() const = 0; diff --git a/third_party/xla/xla/pjrt/pjrt_client_test.cc b/third_party/xla/xla/pjrt/pjrt_client_test.cc index 7f6b62f30afb13..7343158f36e4d8 100644 --- a/third_party/xla/xla/pjrt/pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/pjrt_client_test.cc @@ -799,6 +799,11 @@ ENTRY RuntimeDonationDenialMustAliasFails() -> f32[2, 2] { TEST(PjRtClientTest, GetDefaultLayout) {} +TEST(PjRtClientTest, IsCApiTest) { + TF_ASSERT_OK_AND_ASSIGN(auto client, GetClient()); + EXPECT_FALSE(client->IsCApi()); +} + TEST(PjRtClientTest, ClearPeakMemory) { TF_ASSERT_OK_AND_ASSIGN(auto client, GetClient()); PjRtDevice* device = client->addressable_devices()[0]; From 2289ff834631e227d838f5017225cb1e454f800c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 4 Sep 2026 18:04:23 -0700 Subject: [PATCH 08/12] XLA:PJRT: Optimize transpose loop ordering Improve loop scheduling heuristic in PJRT TransposePlan when dealing with large strides across dimensions. When ordering loops in permuted transposes, model memory traffic based on expected cache-line misses (treating strides >= 128 bytes as full cache-line misses on input and output) to push large DRAM strides into outer loops, using the smaller stride as a tie-breaker. Preserve baseline std::min prioritization for identity permutations (pure tile packing), tile interior loops, and non-permuted dimensions. PiperOrigin-RevId: 976578359 --- third_party/xla/xla/pjrt/transpose.cc | 67 +++++++++++++++++++++------ third_party/xla/xla/pjrt/transpose.h | 1 + 2 files changed, 55 insertions(+), 13 deletions(-) diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc index 41afad3c6099ed..4cf23bec75c33f 100644 --- a/third_party/xla/xla/pjrt/transpose.cc +++ b/third_party/xla/xla/pjrt/transpose.cc @@ -1137,7 +1137,16 @@ absl::Status TransposePlan::Initialize() { ++pos_stride1a_in_a; } + is_identity_permutation_ = true; + for (int i = 0; i < permutation_.size(); ++i) { + if (permutation_[i] != i) { + is_identity_permutation_ = false; + break; + } + } + b_dims_ = Permute(a_dims_, permutation_); + ComputeStrides(elem_size_in_bytes_, b_dims_, b_tiling_, ldb_, ldb_tile_); // Find the innermost dimension of B that is stride 1 element. We know such a @@ -1377,7 +1386,17 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { loop_order.clear(); loop_order.reserve(remaining.size()); - auto soft_cost = [&](const Loop& l) { + // Computes a sort key for each candidate loop, returning a 3-tuple: + // + // Lexicographical ordering selects the smallest tuple first, placing that + // loop into an outer position (from slowest-varying to fastest-varying): + // 1. Contiguity category (non-contiguous loops placed in outer loops). + // 2. Primary cost: cache-line misses (higher miss cost placed in outer + // loops). + // 3. Secondary tie-breaker: when cache misses tie (e.g. both loops have + // strides >= 128B), the loop with larger minimum stride goes outer to + // favor better locality in at least one buffer for the inner loop. + auto soft_cost = [&](const Loop& l) -> std::tuple { int64_t a_stride = std::abs(l.lda); if (!inner_kernel_is_memcpy_ && l.is_inner_dim_in_a) { a_stride *= inner_block_elems_ * outer_block_elems_a_; @@ -1387,21 +1406,41 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { b_stride *= inner_block_elems_ * outer_block_elems_b_; } - double stride; switch (chunk_contiguity_) { case ChunkContiguity::kOutput: - stride = b_stride; - return std::make_tuple(0, -stride); + return std::make_tuple(0, -b_stride, 0.0); case ChunkContiguity::kInput: - stride = a_stride; - return std::make_tuple(0, -stride); - case ChunkContiguity::kNone: - // Add a small penalty to the input strides: given the choice between - // consecutive writes and consecutive reads, we would prefer consecutive - // writes. - constexpr double kPenalty = 1.01; - stride = std::min(a_stride * kPenalty, b_stride); - return std::make_tuple(l.contiguity, -stride); + return std::make_tuple(0, -a_stride, 0.0); + case ChunkContiguity::kNone: { + constexpr double kOutputPenalty = 1.01; + const double min_stride = + std::min(a_stride * kOutputPenalty, b_stride); + + // Pure tile-packing identity linearizers/delinearizers preserve the + // baseline std::min cost to keep 2D hardware tile pairings intact. + if (is_identity_permutation_) { + return std::make_tuple(l.contiguity, -min_stride, 0.0); + } + + // Model the memory traffic of placing this loop in an inner position + // based on the count of cache lines missed (bytes pulled from L2 to + // L1). Memory transfers occur at cache-line granularity (128-byte cache + // line pairs on modern CPUs like AMD Zen 3/4). + // - Strides < 128 bytes share cache lines across iterations (spatial + // reuse). + // - Strides >= 128 bytes incur a full cache-line miss per iteration. + constexpr double kCacheLineBytes = 128.0; + + const double misses_a = std::min(1.0, a_stride / kCacheLineBytes); + const double misses_b = std::min(1.0, b_stride / kCacheLineBytes); + const double base_cost = misses_a + misses_b * kOutputPenalty; + + // Break ties among loops with full cache-line misses using the minimum + // stride across both buffers (favoring loops with better locality in + // at least one buffer, such as tile interior loops over tile exterior + // loops). + return std::make_tuple(l.contiguity, -base_cost, -min_stride); + } } }; @@ -1417,6 +1456,7 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { } // Hard constraint 2: tile ordering. + // A tile interior MUST come after its corresponding tile exterior. if (l.tile_interior) { auto is_exterior = [&](const Loop& r) { return !r.tile_interior && r.dim_in_a == l.dim_in_a; @@ -1434,6 +1474,7 @@ void TransposePlan::ChooseLoopOrder(std::vector& loop_order) const { loop_order.push_back(std::move(remaining[best_idx])); remaining.erase(remaining.begin() + best_idx); } + VLOG(5) << "After loop ordering sort: " << absl::StrJoin(loop_order, ", ", [](std::string* out, const Loop& l) { diff --git a/third_party/xla/xla/pjrt/transpose.h b/third_party/xla/xla/pjrt/transpose.h index 9a91f53a488c76..1b7105dc9b1d60 100644 --- a/third_party/xla/xla/pjrt/transpose.h +++ b/third_party/xla/xla/pjrt/transpose.h @@ -353,6 +353,7 @@ class TransposePlan { absl::InlinedVector b_tiling_; bool a_is_tiled_ = false; bool b_is_tiled_ = false; + bool is_identity_permutation_ = false; // Per-chunk loop nests. Each loop nest has its own start/end bounds // representing one chunk of the work. From 6229c2d456502900f79a36f5aae891e6741dd5bd Mon Sep 17 00:00:00 2001 From: Bhatu Date: Fri, 4 Sep 2026 18:59:44 -0700 Subject: [PATCH 09/12] Clean up MakeFakeArguments API to use FakeArgumentsOptions and migrate all callsites - Introduce FakeArgumentsOptions struct to replace positional parameters for MakeFakeArguments and MakeDataflowConstrainedArguments. - Migrate all callsites to use FakeArgumentsOptions or default options {}. PiperOrigin-RevId: 976593749 --- .../gpu/codegen/triton/dot_algorithms_test.cc | 21 ++-- .../tests/fusion_emitter_device_test.cc | 4 +- .../xla/backends/gpu/tests/ragged_dot_test.cc | 40 +++---- .../backends/gpu/tests/regression_dot_test.cc | 8 +- .../xla/service/gpu/model/hlo_op_profiler.cc | 17 ++- .../hlo_runner_agnostic_reference_mixin.h | 7 +- .../tests/hlo_runner_agnostic_test_base.cc | 34 +++--- third_party/xla/xla/tests/test_utils.cc | 64 ++++------ third_party/xla/xla/tests/test_utils.h | 113 ++++++++---------- third_party/xla/xla/tests/test_utils_test.cc | 59 +++------ .../hlo_bisect/restricted/hlo_bisect_utils.cc | 4 +- .../hlo_isolation_test_base_test.cc | 56 ++++----- .../xla/xla/tools/matmul_perf_table_gen.cc | 21 ++-- .../functional_hlo_runner.cc | 19 ++- third_party/xla/xla/tools/run_hlo_module.cc | 15 ++- 15 files changed, 207 insertions(+), 275 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc index 7c8f2862bea5f1..26997ce8df3b22 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc @@ -1960,12 +1960,10 @@ TEST_P(PrecisionTests, PrecisionCheck) { std::unique_ptr test_module, GetSimpleDotModule(kLhsOuterDim, kRhsOuterDim, kContractingDim, algorithm, backend)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector fake_arguments, - MakeFakeArguments(test_module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/23)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 23; + ASSERT_OK_AND_ASSIGN(std::vector fake_arguments, + MakeFakeArguments(test_module.get(), options)); // Ensure there are no negative arguments to avoid unbounded relative errors // due to subtracting two similarly large numbers. MakeNonNegative(fake_arguments); @@ -2030,13 +2028,10 @@ TEST_P(PrecisionTests, CheckPrecisionDegradationAlongKDimension) { TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr test_module, GetSimpleDotModule(kMSize, kNSize, k, algorithm, backend)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector fake_arguments, - MakeFakeArguments(test_module.get(), /*pseudo_random=*/ - true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/23)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 23; + ASSERT_OK_AND_ASSIGN(std::vector fake_arguments, + MakeFakeArguments(test_module.get(), options)); // Ensure there are no negative arguments to avoid unbounded relative errors // due to subtracting two similarly large numbers. MakeNonNegative(fake_arguments); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc index ba82ec58c82376..11c9f5ed824631 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc @@ -2571,8 +2571,10 @@ class TritonScaledDotTestBase : public TritonEmitterTest { absl::StatusOr> MakeScaledDotArguments( const HloModule* module) { std::minstd_rand0 engine; + FakeArgumentsOptions options; + options.engine = &engine; ABSL_ASSIGN_OR_RETURN(std::vector arguments, - MakeFakeArguments(module, &engine)); + MakeFakeArguments(module, options)); if (arguments.size() != 4) { return absl::InternalError(absl::StrCat( "Expected 4 scaled-dot arguments, got ", arguments.size())); diff --git a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc index cb9450a4362ce9..8cf1eaf1375ebe 100644 --- a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc @@ -44,12 +44,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR1({1, 2, 3}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -71,12 +69,10 @@ TEST_F(RaggedDotTest, NonContractingWithBatchDims) { lhs_ragged_dims={1}, rhs_group_dims={1} })"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=9. fake_arguments[2] = LiteralUtil::CreateR2({{4, 5}, {7, 2}, {6, 3}}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -98,12 +94,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR1({4, 2}); EXPECT_TRUE(RunAndCompare(std::move(module), @@ -125,12 +119,10 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/10)); + FakeArgumentsOptions options; + options.max_bits_of_precision = 10; + TF_ASSERT_OK_AND_ASSIGN(auto fake_arguments, + MakeFakeArguments(module.get(), options)); // Set group sizes to reasonable numbers for ragged_dim_size=6. fake_arguments[2] = LiteralUtil::CreateR2({{1, 2, 3}, {3, 2, 1}}); EXPECT_TRUE(RunAndCompare(std::move(module), diff --git a/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc b/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc index af6e1098a2cb8f..1bbfaa4bdbddf3 100644 --- a/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/regression_dot_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include +#include #include #include "xla/error_spec.h" #include "xla/literal_util.h" @@ -45,11 +46,8 @@ ENTRY main { ROOT R = bf16[3072] reduce(prod, zero), dimensions={0}, to_apply=sum } )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - auto fake_arguments, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + ASSERT_OK_AND_ASSIGN(auto fake_arguments, MakeFakeArguments(module.get())); EXPECT_TRUE(RunAndCompare(std::move(module), LiteralUtil::MakePointers(fake_arguments), diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc index 062cd41a9d6268..bdfef0a45c0dea 100644 --- a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc +++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.cc @@ -271,12 +271,17 @@ absl::StatusOr HloOpProfiler::MeasureOpChainDuration( std::minstd_rand0 engine; // Some operations have dynamic duration that depends on the input values. // Measure each operation with small and large inputs and average. - std::vector args_small = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/false) - .value(); - std::vector args_large = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/true) - .value(); + FakeArgumentsOptions small_options; + small_options.engine = &engine; + small_options.use_large_range = false; + ABSL_ASSIGN_OR_RETURN(std::vector args_small, + MakeFakeArguments(module.get(), small_options)); + + FakeArgumentsOptions large_options; + large_options.engine = &engine; + large_options.use_large_range = true; + ABSL_ASSIGN_OR_RETURN(std::vector args_large, + MakeFakeArguments(module.get(), large_options)); const absl::Time t_compile_start = absl::Now(); ABSL_ASSIGN_OR_RETURN(std::unique_ptr ex, runner_.CreateExecutable(std::move(module), diff --git a/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h b/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h index 98ff2277fa4704..48ea1e3061bbb7 100644 --- a/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h +++ b/third_party/xla/xla/tests/hlo_runner_agnostic_reference_mixin.h @@ -127,11 +127,10 @@ class HloRunnerAgnosticReferenceMixin : public T { const std::function& reference_preprocessor = nullptr, const std::function& test_preprocessor = nullptr, const std::optional args_max_bits_of_precision = std::nullopt) { + FakeArgumentsOptions options; + options.max_bits_of_precision = args_max_bits_of_precision; const absl::StatusOr> fake_arguments = - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - args_max_bits_of_precision); + MakeFakeArguments(module.get(), options); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status().message(); } diff --git a/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc b/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc index 48195abd4cbcb9..ccb91b2b5bc2ff 100644 --- a/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc +++ b/third_party/xla/xla/tests/hlo_runner_agnostic_test_base.cc @@ -295,11 +295,8 @@ HloRunnerAgnosticTestBase::RunAndCompareTwoModulesReplicated( std::unique_ptr module_0, std::unique_ptr module_1, const bool run_hlo_passes, const bool use_threads, const std::optional& error) { - const absl::StatusOr> fake_arguments = MakeFakeArguments( - /*module=*/module_0.get(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt); + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.get()); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -372,9 +369,10 @@ ::testing::AssertionResult HloRunnerAgnosticTestBase::RunAndCompareTwoModules( << absl::StrJoin(mismatches, ", "); } - const absl::StatusOr> fake_arguments = MakeFakeArguments( - module_0.get(), /*pseudo_random=*/true, /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, args_max_bits_of_precision); + FakeArgumentsOptions options; + options.max_bits_of_precision = args_max_bits_of_precision; + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.get(), options); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -483,11 +481,8 @@ HloRunnerAgnosticTestBase::RunAndCompareTwoExecutables( << "Error : mismatching parameter shapes for parameters " << absl::StrJoin(mismatches, ", "); } - absl::StatusOr> fake_arguments = MakeFakeArguments( - /*module=*/module_0.value(), /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt); + absl::StatusOr> fake_arguments = + MakeFakeArguments(module_0.value()); if (!fake_arguments.ok()) { return ::testing::AssertionFailure() << fake_arguments.status(); } @@ -516,11 +511,18 @@ ::testing::AssertionResult HloRunnerAgnosticTestBase::Run( << "Error while parsing HLO text format: " << module.status().ToString(); } - const std::vector fake_arguments = - MakeFakeArguments(module->get(), use_random_data).value(); + FakeArgumentsOptions options; + options.pseudo_random = use_random_data; + const absl::StatusOr> fake_arguments = + MakeFakeArguments(module->get(), options); + if (!fake_arguments.ok()) { + return ::testing::AssertionFailure() + << "Error while generating fake arguments: " + << fake_arguments.status().ToString(); + } std::vector fake_argument_ptrs; absl::c_transform( - fake_arguments, std::back_inserter(fake_argument_ptrs), + *fake_arguments, std::back_inserter(fake_argument_ptrs), [](const Literal& literal) { return const_cast(&literal); }); if (backend_config) { diff --git a/third_party/xla/xla/tests/test_utils.cc b/third_party/xla/xla/tests/test_utils.cc index a34c66b735cb38..c99759665732fd 100644 --- a/third_party/xla/xla/tests/test_utils.cc +++ b/third_party/xla/xla/tests/test_utils.cc @@ -338,36 +338,16 @@ absl::StatusOr MakeConstrainedArgument( } // namespace absl::StatusOr> MakeFakeArguments( - const HloModule* module, bool pseudo_random, bool use_large_range, - bool treat_gte_as_data_formatting, - std::optional max_bits_of_precision, std::minstd_rand0* engine, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { - if (!pseudo_random) { - return MakeFakeArguments(module, nullptr, use_large_range, - treat_gte_as_data_formatting, - max_bits_of_precision, generate_aligned_ds_indices, - get_index_known_zeroes); - } - if (engine == nullptr) { - auto new_engine = - pseudo_random ? std::make_unique() : nullptr; - return MakeFakeArguments(module, new_engine.get(), use_large_range, - treat_gte_as_data_formatting, - max_bits_of_precision, generate_aligned_ds_indices, - get_index_known_zeroes); + const HloModule* module, const FakeArgumentsOptions& options) { + std::unique_ptr default_engine; + std::minstd_rand0* engine = options.engine; + if (!options.pseudo_random) { + engine = nullptr; + } else if (engine == nullptr) { + default_engine = std::make_unique(); + engine = default_engine.get(); } - return MakeFakeArguments(module, engine, use_large_range, - treat_gte_as_data_formatting, max_bits_of_precision, - generate_aligned_ds_indices, get_index_known_zeroes); -} -absl::StatusOr> MakeFakeArguments( - const HloModule* module, std::minstd_rand0* engine, bool use_large_range, - bool treat_gte_as_data_formatting, - std::optional max_bits_of_precision, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { ABSL_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(*module)); const auto params = module->entry_computation()->parameter_instructions(); std::vector arguments(params.size()); @@ -386,33 +366,35 @@ absl::StatusOr> MakeFakeArguments( ABSL_ASSIGN_OR_RETURN( arguments[i], MakeConstrainedArgument( - *dataflow, *params[i], param_shape, engine, use_large_range, - treat_gte_as_data_formatting, max_bits_of_precision, - generate_aligned_ds_indices, get_index_known_zeroes)); + *dataflow, *params[i], param_shape, engine, options.use_large_range, + options.treat_gte_as_data_formatting, options.max_bits_of_precision, + options.generate_aligned_ds_indices, + options.get_index_known_zeroes)); } return std::move(arguments); } absl::StatusOr> MakeDataflowConstrainedArguments( - const HloModule* module, std::minstd_rand0* engine, bool use_large_range, - std::optional max_bits_of_precision, - bool generate_aligned_ds_indices, - GetIndexKnownZeroesFn get_index_known_zeroes) { + const HloModule* module, const FakeArgumentsOptions& options) { std::unique_ptr default_engine; - if (engine == nullptr) { + std::minstd_rand0* engine = options.engine; + if (!options.pseudo_random) { + engine = nullptr; + } else if (engine == nullptr) { default_engine = std::make_unique(); engine = default_engine.get(); } - ABSL_ASSIGN_OR_RETURN(auto constraint_states, - ConstraintPropagator::Run(*module, get_index_known_zeroes)); + ABSL_ASSIGN_OR_RETURN( + auto constraint_states, + ConstraintPropagator::Run(*module, options.get_index_known_zeroes)); auto make_literal_for_state = [&](const Shape& shape, const ConstraintState& state, absl::string_view target_name) -> absl::StatusOr { ConstraintInterval interval = state.GetConstraintInterval(); StructuralConstraints structure = state.GetStructuralConstraints(); - if (!generate_aligned_ds_indices) { + if (!options.generate_aligned_ds_indices) { structure.alignment = std::nullopt; } std::optional> limit = std::nullopt; @@ -454,8 +436,8 @@ absl::StatusOr> MakeDataflowConstrainedArguments( limit = {min_val, max_val}; } return MakeFakeLiteral(shape, engine, limit, structure.needs_sorted_indices, - structure.no_duplicates, use_large_range, - max_bits_of_precision, structure.alignment, + structure.no_duplicates, options.use_large_range, + options.max_bits_of_precision, structure.alignment, structure.known_zeroes_mask, /*float_generator=*/nullptr, interval); }; diff --git a/third_party/xla/xla/tests/test_utils.h b/third_party/xla/xla/tests/test_utils.h index 23976222b8a562..da152feeb74adf 100644 --- a/third_party/xla/xla/tests/test_utils.h +++ b/third_party/xla/xla/tests/test_utils.h @@ -65,6 +65,53 @@ class PseudorandomGenerator { using GetIndexKnownZeroesFn = std::function(const HloInstruction*, int64_t)>; +// Options for generating fake arguments with MakeFakeArguments and +// MakeDataflowConstrainedArguments. +struct FakeArgumentsOptions { + // Optional random number generator. Passing a generator enables generation + // of different random values across sequential calls by reusing the same + // engine. + std::minstd_rand0* engine = nullptr; + + // If pseudo_random is true, the generated numbers will be generated + // deterministically in a pseudo random way unless the values are constrained + // to be e.g. init values as above. If pseudo_random is false, the returned + // values will be generated in a faster way that yields less interesting data, + // e.g. the values may all be just the same value. + // + // TODO(b/79942829): Make interesting argument generation fast enough that + // using pseudo_random does not save any noticeable amount of time so that the + // parameter can be removed. + bool pseudo_random = true; + + // If use_large_range is false, the generated floating point numbers will be + // sampled from a small range of possible values. If use_large_range is true, + // the generated floating point numbers will be sampled from a uniform-log + // distribution of most possible floats, with a small chance to instead be + // sampled from a list of special floating point values (such as 0, inf, + // etc.). + bool use_large_range = false; + + // If treat_gte_as_data_formatting is true, GetTupleElement instructions are + // treated as data formatting operations when tracking parameter constraints, + // allowing constraints to propagate through tuple deconstruction. + bool treat_gte_as_data_formatting = false; + + // If max_bits_of_precision is set to a number, then floating point & integer + // types will be constrained to be represented in that number of bits. Setting + // it to 5 for integers would mean it only creates integers between -32 and + // 32. + std::optional max_bits_of_precision = std::nullopt; + + // If `generate_aligned_ds_indices` is true, the generated indices will be + // aligned to the given alignment. + bool generate_aligned_ds_indices = false; + + // If `get_index_known_zeroes` is set, the generated indices will have the + // given number of zeroes in the given dimension. + GetIndexKnownZeroesFn get_index_known_zeroes = nullptr; +}; + // Generates a vector of arguments containing fake data. The number, shape and // layout of the arguments is appropriate for given HLO module. // @@ -79,80 +126,20 @@ using GetIndexKnownZeroesFn = // (3) Keys of key/value sorts should contain no duplicates. // // These constraints are best-effort only. -// -// If max_bits_of_precision is set to a number, then floating point & integer -// types will be constrained to be represented in that number of bits. Setting -// it to 5 for integers would mean it only creates integers between -32 and 32. -// -// If pseudo_random is true, the generated numbers will be generated -// deterministically in a pseudo random way unless the values are constrated to -// be e.g. init values as above. If pseudo_random is false, the returned values -// will be generated in a faster way that yields less interesting data, e.g. the -// values may all be just the same value. -// -// If use_large_range is false, the generated floating point numbers will be -// sampled from a small range of possible values. If use_large_range is true, -// the generated floating point numbers will be sampled from a uniform-log -// distribution of most possible floats, with a small chance to instead be -// sampled from a list of special floating point values (such as 0, inf, etc.). -// -// TODO(b/79942829): Make interesting argument generation fast enough that using -// pseudo_random does not save any noticeable amount of time so that the -// parameter can be removed. -// -// If `generate_aligned_ds_indices` is true, the generated indices will be -// aligned to the given alignment. If `get_index_known_zeroes` is set, the -// generated indices will have the given number of zeroes in the given -// dimension. -absl::StatusOr> MakeFakeArguments( - const HloModule* module, bool pseudo_random = true, - bool use_large_range = false, bool treat_gte_as_data_formatting = false, - std::optional max_bits_of_precision = std::nullopt, - std::minstd_rand0* engine = nullptr, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); - -// Overload which accepts a random number generator. This enables generation of -// different random values with sequential calls to MakeFakeArguments by reusing -// the same generator. absl::StatusOr> MakeFakeArguments( - const HloModule* module, std::minstd_rand0* engine, - bool use_large_range = false, bool treat_gte_as_data_formatting = false, - std::optional max_bits_of_precision = std::nullopt, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); + const HloModule* module, const FakeArgumentsOptions& options = {}); // Generates a vector of arguments containing fake data using reverse constraint // propagation. The constraint propagator seeds initial constraints based on HLO // op semantics (e.g., `sqrt(x)` implies `x >= 0`) and then propagates these // constraints backward through the graph. This allows generating test inputs // that are more likely to be valid for the graph. -// -// If `use_large_range` is false, the generated floating point numbers will be -// sampled from a small range of possible values. If `use_large_range` is true, -// the generated floating point numbers will be sampled from a uniform-log -// distribution of most possible floats, with a small chance to instead be -// sampled from a list of special floating point values (such as 0, inf, etc.). -// -// If `max_bits_of_precision` is set to a number, then floating point & integer -// types will be constrained to be represented in that number of bits. Setting -// it to 5 for integers would mean it only creates integers between -32 and 32. -// -// If `generate_aligned_ds_indices` is true, the generated indices will be -// aligned to the given alignment. -// -// If `get_index_known_zeroes` is set, the generated indices will have the given -// number of zeroes in the given dimension. absl::StatusOr> MakeDataflowConstrainedArguments( - const HloModule* module, std::minstd_rand0* engine = nullptr, - bool use_large_range = false, - std::optional max_bits_of_precision = std::nullopt, - bool generate_aligned_ds_indices = false, - GetIndexKnownZeroesFn get_index_known_zeroes = nullptr); + const HloModule* module, const FakeArgumentsOptions& options = {}); // Check that a given module satisfies various constraints before trying to // execute it. -absl::Status VerifyHloModule(HloModule* const module, bool layout_sensitive, +absl::Status VerifyHloModule(HloModule* module, bool layout_sensitive, bool allow_mixed_precision); // Creates a dot op with operands 'lhs' and 'rhs' that contracts dimension 1 of diff --git a/third_party/xla/xla/tests/test_utils_test.cc b/third_party/xla/xla/tests/test_utils_test.cc index 90192dcb520d96..20d7de9c4c62aa 100644 --- a/third_party/xla/xla/tests/test_utils_test.cc +++ b/third_party/xla/xla/tests/test_utils_test.cc @@ -258,11 +258,11 @@ ENTRY cluster_13361217111314620287__.11 { )") .value(); - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), /*pseudo_random=*/true, - /*use_large_range=*/true, - /*treat_gte_as_data_formatting=*/true)); + FakeArgumentsOptions options; + options.use_large_range = true; + options.treat_gte_as_data_formatting = true; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 1); const Shape& indices_shape = args[0].shape().tuple_shapes()[0]; @@ -474,16 +474,10 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s32[]) -> s8[131072,2048] { return std::nullopt; }; - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + FakeArgumentsOptions options; + options.get_index_known_zeroes = index_known_zeroes_fn; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 2); int32_t index = args[1].Get({}); @@ -515,16 +509,10 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) return std::nullopt; }; - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + FakeArgumentsOptions options; + options.get_index_known_zeroes = index_known_zeroes_fn; + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 3); int32_t index = args[2].Get({}); @@ -533,12 +521,7 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) TF_ASSERT_OK_AND_ASSIGN( std::vector args2, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*generate_aligned_ds_indices=*/false, - index_known_zeroes_fn)); + MakeDataflowConstrainedArguments(module.get(), options)); ASSERT_EQ(args2.size(), 3); int32_t index2 = args2[2].Get({}); EXPECT_EQ(index2 & index_known_bits_zero, 0); @@ -559,11 +542,8 @@ ENTRY main { } )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - TF_ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false)); + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeDataflowConstrainedArguments(module.get())); ASSERT_EQ(args.size(), 1); args[0].EachCell([](absl::Span indices, float value) { EXPECT_GT(value, 0.0f); @@ -580,11 +560,8 @@ ENTRY main { } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeDataflowConstrainedArguments(module.get(), - /*engine=*/nullptr, - /*use_large_range=*/false)); + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeDataflowConstrainedArguments(module.get())); ASSERT_EQ(args.size(), 1); args[0].EachCell([](absl::Span indices, int32_t value) { EXPECT_GE(value, 1); }); diff --git a/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc b/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc index 7028adf7389728..78d82ddb469450 100644 --- a/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc +++ b/third_party/xla/xla/tools/hlo_bisect/restricted/hlo_bisect_utils.cc @@ -194,8 +194,10 @@ MiscompareChecker::MiscompareChecker(HloModule* module, // Generate input data and store the data for all the execution. std::minstd_rand0 rng_engine; if (input_data.empty()) { + FakeArgumentsOptions options; + options.engine = &rng_engine; absl::StatusOr> input_status = - MakeFakeArguments(module, &rng_engine); + MakeFakeArguments(module, options); CHECK(input_status.ok()); input_data_ = std::move(input_status).value(); } else { diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc index ecd7b78a58a240..1fd28ce0b54187 100644 --- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc +++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_test_base_test.cc @@ -510,23 +510,17 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s32[]) -> s8[131072,2048] { } )")); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - [](const HloInstruction* use, - int64_t sliced_dim) -> std::optional { - if (use->opcode() == HloOpcode::kDynamicSlice && - sliced_dim == 0) { - return 131071; - } - return std::nullopt; - })); + FakeArgumentsOptions options; + options.get_index_known_zeroes = + [](const HloInstruction* use, + int64_t sliced_dim) -> std::optional { + if (use->opcode() == HloOpcode::kDynamicSlice && sliced_dim == 0) { + return 131071; + } + return std::nullopt; + }; + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options)); ASSERT_EQ(args.size(), 2); int32_t index = args[1].Get({}); @@ -548,23 +542,17 @@ ENTRY %main (param_1: s8[262144,2048], param_2: s8[131072,2048], param_3: s32[]) } )")); - ASSERT_OK_AND_ASSIGN( - std::vector args, - MakeFakeArguments(module.get(), - /*pseudo_random=*/true, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, - /*engine=*/nullptr, - /*generate_aligned_ds_indices=*/false, - [](const HloInstruction* use, - int64_t sliced_dim) -> std::optional { - if (use->opcode() == HloOpcode::kDynamicUpdateSlice && - sliced_dim == 0) { - return 131071; - } - return std::nullopt; - })); + FakeArgumentsOptions options_dus; + options_dus.get_index_known_zeroes = + [](const HloInstruction* use, + int64_t sliced_dim) -> std::optional { + if (use->opcode() == HloOpcode::kDynamicUpdateSlice && sliced_dim == 0) { + return 131071; + } + return std::nullopt; + }; + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get(), options_dus)); ASSERT_EQ(args.size(), 3); int32_t index = args[2].Get({}); diff --git a/third_party/xla/xla/tools/matmul_perf_table_gen.cc b/third_party/xla/xla/tools/matmul_perf_table_gen.cc index ac16e18c9df78b..530c137195fae0 100644 --- a/third_party/xla/xla/tools/matmul_perf_table_gen.cc +++ b/third_party/xla/xla/tools/matmul_perf_table_gen.cc @@ -402,23 +402,28 @@ absl::Duration MatmulPerfTableGen::Profile(std::unique_ptr module) { // Flip flop between arguments to prevent caching. std::minstd_rand0 engine; - std::vector args_small = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/false) - .value(); - std::vector args_large = MakeFakeArguments(module.get(), &engine, - /*use_large_range=*/true) - .value(); + FakeArgumentsOptions small_options; + small_options.engine = &engine; + small_options.use_large_range = false; + auto args_small = MakeFakeArguments(module.get(), small_options); + CHECK_OK(args_small); + + FakeArgumentsOptions large_options; + large_options.engine = &engine; + large_options.use_large_range = true; + auto args_large = MakeFakeArguments(module.get(), large_options); + CHECK_OK(args_large); std::unique_ptr compiled = Compile(std::move(module)); // First run to warm up stuff. - CHECK_OK(runner_.ExecuteWithExecutable(compiled.get(), args_small).status()); + CHECK_OK(runner_.ExecuteWithExecutable(compiled.get(), *args_small).status()); // Trace `kNumProfilingRuns` times to get decent measurement. std::unique_ptr tracer = HloOpProfiler::GetKernelTracer(); for (int i = 0; i < kNumProfilingRuns; i++) { - Measure(runner_, compiled.get(), args_small, args_large); + Measure(runner_, compiled.get(), *args_small, *args_large); } return absl::Nanoseconds(std::move(*tracer).getMedianKernelTimeNs()); diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc index 9da2396d55903f..8d1ab3ec5e91aa 100644 --- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc +++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc @@ -1004,23 +1004,18 @@ CreateArgumentsOnDevice(PjRtClient& client, } } } else { + FakeArgumentsOptions options; + options.engine = engine; + options.pseudo_random = kUseRandomInputs; if (flatten_arguments) { - ABSL_ASSIGN_OR_RETURN( - LiteralVec tupled_argument_literals, - MakeFakeArguments(my_hlo_module, kUseRandomInputs, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, engine)); + ABSL_ASSIGN_OR_RETURN(LiteralVec tupled_argument_literals, + MakeFakeArguments(my_hlo_module, options)); CHECK_EQ(tupled_argument_literals.size(), 1); CHECK(tupled_argument_literals.front().shape().IsTuple()); argument_literals = tupled_argument_literals.front().DecomposeTuple(); } else { - ABSL_ASSIGN_OR_RETURN( - argument_literals, - MakeFakeArguments(my_hlo_module, kUseRandomInputs, - /*use_large_range=*/false, - /*treat_gte_as_data_formatting=*/false, - /*max_bits_of_precision=*/std::nullopt, engine)); + ABSL_ASSIGN_OR_RETURN(argument_literals, + MakeFakeArguments(my_hlo_module, options)); } if (kUseSharedInputs) { break; diff --git a/third_party/xla/xla/tools/run_hlo_module.cc b/third_party/xla/xla/tools/run_hlo_module.cc index f219c1cbc51884..ffca6b41055396 100644 --- a/third_party/xla/xla/tools/run_hlo_module.cc +++ b/third_party/xla/xla/tools/run_hlo_module.cc @@ -242,12 +242,15 @@ absl::Status RunAndCompareInternal( .status()); } - ABSL_ASSIGN_OR_RETURN(auto args, - copy_result_on_failure( - MakeFakeArguments(test_module.get(), engine, - options.use_large_float_range, - options.treat_gte_as_data_formatting), - ModuleResult::kOtherError, test_run_result)); + FakeArgumentsOptions fake_arguments_options; + fake_arguments_options.engine = engine; + fake_arguments_options.use_large_range = options.use_large_float_range; + fake_arguments_options.treat_gte_as_data_formatting = + options.treat_gte_as_data_formatting; + ABSL_ASSIGN_OR_RETURN(auto args, copy_result_on_failure( + MakeFakeArguments(test_module.get(), + fake_arguments_options), + ModuleResult::kOtherError, test_run_result)); // Use provided input literals as arguments, if any. if (iteration_literals_proto != nullptr && iteration_literals_proto->arguments_size() != 0) { From 4357817bc0ac19ed49bfb1d22c472c365ab0f89b Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Fri, 4 Sep 2026 19:17:34 -0700 Subject: [PATCH 10/12] Add helpers for MLIR import and export APIs in Python. PiperOrigin-RevId: 976598293 --- tensorflow/compiler/mlir/lite/BUILD | 1 + tensorflow/compiler/mlir/lite/python/BUILD | 14 ++ .../lite/python/_pywrap_converter_api.pyi | 3 +- .../mlir/lite/python/converter_python_api.cc | 18 ++- .../mlir/lite/python/converter_python_api.h | 13 +- .../python/converter_python_api_wrapper.cc | 32 +++- .../mlir/lite/python/flatbuffer_to_mlir.cc | 143 +++++++++++++++++- .../mlir/lite/python/flatbuffer_to_mlir.h | 18 ++- .../mlir/lite/python/wrap_converter.py | 34 ++++- tensorflow/lite/python/BUILD | 2 +- tensorflow/lite/python/analyzer.py | 21 ++- tensorflow/lite/python/convert.py | 61 ++++++++ tensorflow/lite/python/lite.py | 3 + tensorflow/lite/python/lite_test.py | 38 +++++ 14 files changed, 373 insertions(+), 28 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index 435b63a90d65b5..a2acd13d151f4a 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -47,6 +47,7 @@ package_group( "//third_party/iree/...", "//third_party/odml/infra/...", "//third_party/odml/litert/...", + "//third_party/py/ai_edge_jax/...", "//waymo/accelerator/alpine/tools/...", "//waymo/ml/compiler/mlir/...", ], diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 5ed120ebbeade1..5b70d4b2bb3362 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -299,12 +299,26 @@ cc_library( "flatbuffer_to_mlir.h", ], deps = [ + "//tensorflow/compiler/mlir:op_or_arg_name_mapper", + "//tensorflow/compiler/mlir/lite:flatbuffer_export", "//tensorflow/compiler/mlir/lite:flatbuffer_import", + "//tensorflow/compiler/mlir/lite:tensorflow_lite", + "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps", + "//tensorflow/compiler/mlir/tensorflow", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:BytecodeWriter", + "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Parser", + "@llvm-project//mlir:QuantOps", "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", + "@stablehlo//:stablehlo_ops", + "@stablehlo//:vhlo_ops", ], ) diff --git a/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi b/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi index 401c020590f898..8125d0b9c17a4a 100644 --- a/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi +++ b/tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi @@ -17,6 +17,7 @@ def Convert(model_flags_proto_txt_raw: object, converter_flags_proto_txt_raw: ob def ConvertMlirBytecode(converter_flags_proto_txt_raw: object, model_dir_txt_raw: object, output_file_path_raw: object) -> object: ... def ExperimentalMlirQuantizeModel(input_contents_txt_raw: object, disable_per_channel: bool = ..., fully_quantize: bool = ..., inference_type: int = ..., input_data_type: int = ..., output_data_type: int = ..., enable_numeric_verify: bool = ..., enable_whole_model_verify: bool = ..., op_blocklist: object = ..., node_blocklist: object = ..., enable_variable_quantization: bool = ..., disable_per_channel_for_dense_layers: bool = ..., debug_options_proto_txt_raw: object = ...) -> object: ... def ExperimentalMlirSparsifyModel(input_contents_txt_raw: object) -> object: ... -def FlatBufferToMlir(arg0: str, arg1: bool) -> str: ... +def FlatBufferToMlir(model: str, input_is_filepath: bool = ..., bytecode: bool = ..., cl_options: list[str] = ...) -> object: ... +def MlirToFlatBuffer(mlir: str, input_is_filepath: bool = ..., emit_builtin_tflite_ops: bool = ..., emit_select_tf_ops: bool = ..., emit_custom_ops: bool = ..., emit_stablehlo_ops: bool = ...) -> bytes: ... def RegisterCustomOpdefs(custom_opdefs_txt_raw: object) -> object: ... def RetrieveCollectedErrors() -> list: ... diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api.cc b/tensorflow/compiler/mlir/lite/python/converter_python_api.cc index 00f18729a4ec48..a653a6d30de321 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api.cc +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api.cc @@ -16,10 +16,10 @@ limitations under the License. #include -#include #include #include #include +#include #include #include "absl/container/flat_hash_set.h" @@ -517,8 +517,20 @@ std::vector RetrieveCollectedErrors() { } std::string FlatBufferFileToMlir(const std::string& model, - bool input_is_filepath) { - return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath); + bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath, bytecode, + cl_options); +} + +std::string MlirToFlatBufferFile(const std::string& mlir, + bool input_is_filepath, + bool emit_builtin_tflite_ops, + bool emit_select_tf_ops, bool emit_custom_ops, + bool emit_stablehlo_ops) { + return ::tensorflow::MlirToFlatBufferFile( + mlir, input_is_filepath, emit_builtin_tflite_ops, emit_select_tf_ops, + emit_custom_ops, emit_stablehlo_ops); } PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api.h b/tensorflow/compiler/mlir/lite/python/converter_python_api.h index 4ed970e9df0cb9..e53467d3e4586a 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api.h +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api.h @@ -64,8 +64,17 @@ PyObject* RegisterCustomOpdefs(PyObject* list); std::vector RetrieveCollectedErrors(); // Returns MLIR string dump of the given Flatbuffer model. -std::string FlatBufferFileToMlir(const std::string& model, - bool input_is_filepath); +std::string FlatBufferFileToMlir( + const std::string& model, bool input_is_filepath, bool bytecode = false, + const std::vector& cl_options = {}); + +// Converts MLIR (text or bytecode) to a TFLite Flatbuffer. +std::string MlirToFlatBufferFile(const std::string& mlir, + bool input_is_filepath, + bool emit_builtin_tflite_ops = true, + bool emit_select_tf_ops = false, + bool emit_custom_ops = true, + bool emit_stablehlo_ops = false); // Convert slim model to TfLite flatbuffer streamed directly to a file. PyObject* ConvertMlirBytecode(PyObject* converter_flags_proto_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc b/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc index 16af603063a0f4..57198b7ca87339 100644 --- a/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc +++ b/tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include "pybind11/pybind11.h" // from @pybind11 +#include "pybind11/stl.h" // from @pybind11 #include "tensorflow/compiler/mlir/lite/python/converter_python_api.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h" #include "tensorflow/python/lib/core/pybind11_lib.h" @@ -119,12 +120,39 @@ PYBIND11_MODULE(_pywrap_converter_api, m, py::mod_gil_not_used()) { )pbdoc"); m.def( "FlatBufferToMlir", - [](const std::string& model, bool input_is_filepath) { - return tflite::FlatBufferFileToMlir(model, input_is_filepath); + [](const std::string& model, bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + std::string res = tflite::FlatBufferFileToMlir(model, input_is_filepath, + bytecode, cl_options); + if (bytecode) { + return py::object(py::bytes(res)); + } else { + return py::object(py::str(res)); + } }, + py::arg("model"), py::arg("input_is_filepath") = false, + py::arg("bytecode") = false, + py::arg("cl_options") = std::vector(), R"pbdoc( Returns MLIR dump of the given TFLite model. )pbdoc"); + m.def( + "MlirToFlatBuffer", + [](const std::string& mlir, bool input_is_filepath, + bool emit_builtin_tflite_ops, bool emit_select_tf_ops, + bool emit_custom_ops, bool emit_stablehlo_ops) { + std::string res = tflite::MlirToFlatBufferFile( + mlir, input_is_filepath, emit_builtin_tflite_ops, + emit_select_tf_ops, emit_custom_ops, emit_stablehlo_ops); + return py::bytes(res); + }, + py::arg("mlir"), py::arg("input_is_filepath") = false, + py::arg("emit_builtin_tflite_ops") = true, + py::arg("emit_select_tf_ops") = false, py::arg("emit_custom_ops") = true, + py::arg("emit_stablehlo_ops") = false, + R"pbdoc( + Converts MLIR (text or bytecode) into a TFLite FlatBuffer binary tensor. + )pbdoc"); m.def( "ConvertMlirBytecode", [](py::object converter_flags_proto_txt_raw, py::object model_dir_txt_raw, diff --git a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc index b880df7f74a3ca..e80b33d18fba3f 100644 --- a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc +++ b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc @@ -18,21 +18,40 @@ limitations under the License. #include #include +#include "absl/base/attributes.h" +#include "absl/base/const_init.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "llvm/Support/CommandLine.h" #include "llvm/Support/LogicalResult.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/SMLoc.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/raw_ostream.h" +#include "mlir/Bytecode/BytecodeWriter.h" // from @llvm-project +#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/Diagnostics.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/IR/Verifier.h" // from @llvm-project +#include "mlir/Parser/Parser.h" // from @llvm-project #include "mlir/Support/FileUtilities.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo +#include "stablehlo/dialect/VhloOps.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h" #include "tensorflow/compiler/mlir/lite/flatbuffer_import.h" +#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" +#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h" +#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h" +#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h" namespace tensorflow { namespace { @@ -53,8 +72,33 @@ static mlir::OwningOpRef FlatBufferFileToMlirTranslation( } // namespace std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, - bool input_is_filepath) { - // referred logic from mlir::mlirTranslateMain(). + bool input_is_filepath, bool bytecode, + const std::vector& cl_options) { + ABSL_CONST_INIT static absl::Mutex cl_mutex(absl::kConstInit); + absl::MutexLock lock(&cl_mutex); + + // Reset options from any previous invocation. + llvm::cl::ResetAllOptionOccurrences(); + mlir::registerAsmPrinterCLOptions(); + + if (!cl_options.empty()) { + std::vector argv; + argv.reserve(cl_options.size() + 1); + argv.push_back("flatbuffer_to_mlir"); + for (const auto& opt : cl_options) { + argv.push_back(opt.c_str()); + } + std::string cl_errors; + llvm::raw_string_ostream cl_err_stream(cl_errors); + if (!llvm::cl::ParseCommandLineOptions( + argv.size(), argv.data(), "flatbuffer_to_mlir", &cl_err_stream)) { + cl_err_stream.flush(); + if (!cl_errors.empty()) { + llvm::errs() << "Failed to parse MLIR options: " << cl_errors << "\n"; + } + return ""; + } + } std::string errorMessage; std::unique_ptr input; @@ -78,17 +122,102 @@ std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, llvm::SourceMgr sourceMgr; sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc()); + std::string diagnostic_str; + llvm::raw_string_ostream diag_os(diagnostic_str); + mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os); + mlir::OwningOpRef module = FlatBufferFileToMlirTranslation(&sourceMgr, &context); - if (!module || failed(verify(*module))) return ""; + if (!module || failed(verify(*module))) { + diag_os.flush(); + if (!diagnostic_str.empty()) { + llvm::errs() << diagnostic_str << "\n"; + } + return ""; + } std::string mlir_output; llvm::raw_string_ostream output_stream(mlir_output); - // Dump MLIR with eliding large elements. - module->print( - output_stream, - mlir::OpPrintingFlags().useLocalScope().elideLargeElementsAttrs()); + if (bytecode) { + if (mlir::failed(mlir::writeBytecodeToFile(*module, output_stream))) { + llvm::errs() << "Failed to write MLIR bytecode.\n"; + return ""; + } + } else { + mlir::OpPrintingFlags flags; + module->print(output_stream, flags); + } + output_stream.flush(); return mlir_output; } +std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer, + bool input_is_filepath, + bool emit_builtin_tflite_ops, + bool emit_select_tf_ops, bool emit_custom_ops, + bool emit_stablehlo_ops) { + std::string errorMessage; + std::unique_ptr input; + if (input_is_filepath) { + input = mlir::openInputFile(mlir_file_or_buffer, &errorMessage); + if (!input) { + llvm::errs() << errorMessage << "\n"; + return ""; + } + } else { + input = + llvm::MemoryBuffer::getMemBuffer(mlir_file_or_buffer, "mlir", false); + if (!input) { + llvm::errs() << "Can't get llvm::MemoryBuffer\n"; + return ""; + } + } + + mlir::DialectRegistry registry; + registry.insert(); + mlir::RegisterAllTensorFlowDialects(registry); + + mlir::MLIRContext context(registry); + context.printOpOnDiagnostic(true); + + llvm::SourceMgr sourceMgr; + sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc()); + + std::string diagnostic_str; + llvm::raw_string_ostream diag_os(diagnostic_str); + mlir::SourceMgrDiagnosticHandler diag_handler(sourceMgr, &context, diag_os); + + mlir::OwningOpRef module = + mlir::parseSourceFile(sourceMgr, &context); + if (!module || failed(verify(*module))) { + diag_os.flush(); + if (!diagnostic_str.empty()) { + llvm::errs() << diagnostic_str << "\n"; + } else { + llvm::errs() << "Failed to parse MLIR source.\n"; + } + return ""; + } + + std::string serialized_flatbuffer; + tensorflow::OpOrArgLocNameMapper op_or_arg_name_mapper; + tflite::FlatbufferExportOptions options; + options.converter_flags.set_force_select_tf_ops(!emit_builtin_tflite_ops); + options.converter_flags.set_enable_select_tf_ops(emit_select_tf_ops); + options.converter_flags.set_allow_custom_ops(emit_custom_ops); + options.converter_flags.set_use_buffer_offset(true); + options.op_or_arg_name_mapper = &op_or_arg_name_mapper; + + if (!tflite::MlirToFlatBufferTranslateFunction( + *module, options, &serialized_flatbuffer, emit_stablehlo_ops)) { + llvm::errs() << "MlirToFlatBufferTranslateFunction failed.\n"; + return ""; + } + return serialized_flatbuffer; +} + } // namespace tensorflow diff --git a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h index 3164265f9764af..3a297d28c8328a 100644 --- a/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h +++ b/tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h @@ -17,14 +17,24 @@ limitations under the License. #define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_ #include +#include namespace tensorflow { // Translates the given FlatBuffer filename or buffer into MLIR and returns -// translated MLIR as string. -std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer, - bool input_is_filepath); - +// translated MLIR as string or bytecode. +std::string FlatBufferFileToMlir( + const std::string& model_file_or_buffer, bool input_is_filepath, + bool bytecode = false, const std::vector& cl_options = {}); + +// Translates the given MLIR filename or buffer into a TFLite FlatBuffer +// binary string. +std::string MlirToFlatBufferFile(const std::string& mlir_file_or_buffer, + bool input_is_filepath, + bool emit_builtin_tflite_ops = true, + bool emit_select_tf_ops = false, + bool emit_custom_ops = true, + bool emit_stablehlo_ops = false); } // namespace tensorflow #endif // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_ diff --git a/tensorflow/compiler/mlir/lite/python/wrap_converter.py b/tensorflow/compiler/mlir/lite/python/wrap_converter.py index 47689da82d4d40..091be86c08cc73 100644 --- a/tensorflow/compiler/mlir/lite/python/wrap_converter.py +++ b/tensorflow/compiler/mlir/lite/python/wrap_converter.py @@ -98,6 +98,34 @@ def wrapped_retrieve_collected_errors(): return _pywrap_converter_api.RetrieveCollectedErrors() -def wrapped_flat_buffer_file_to_mlir(model, input_is_filepath): - """Wraps FlatBufferFileToMlir with lazy loader.""" - return _pywrap_converter_api.FlatBufferToMlir(model, input_is_filepath) +def wrapped_flat_buffer_file_to_mlir( + model, + input_is_filepath=False, + bytecode=False, + cl_options=None, +): + """Wraps FlatBufferToMlir with lazy loader.""" + if cl_options is None: + cl_options = [] + return _pywrap_converter_api.FlatBufferToMlir( + model, input_is_filepath, bytecode, cl_options + ) + + +def wrapped_mlir_to_flat_buffer( + mlir, + input_is_filepath=False, + emit_builtin_tflite_ops=True, + emit_select_tf_ops=False, + emit_custom_ops=True, + emit_stablehlo_ops=False, +): + """Wraps MlirToFlatBuffer with lazy loader.""" + return _pywrap_converter_api.MlirToFlatBuffer( + mlir, + input_is_filepath, + emit_builtin_tflite_ops, + emit_select_tf_ops, + emit_custom_ops, + emit_stablehlo_ops, + ) diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD index a9cf236970aac8..fc105c7c6fecec 100644 --- a/tensorflow/lite/python/BUILD +++ b/tensorflow/lite/python/BUILD @@ -284,7 +284,7 @@ py_test( ":util", #internal proto upb dep "//third_party/py/numpy", - "//tensorflow:tensorflow_py", + "//tensorflow:tensorflow_py_no_contrib", "//tensorflow/python/client:session", "//tensorflow/python/eager:context", "//tensorflow/python/eager:def_function", diff --git a/tensorflow/lite/python/analyzer.py b/tensorflow/lite/python/analyzer.py index 110d8014786b68..a147508ebe2e08 100644 --- a/tensorflow/lite/python/analyzer.py +++ b/tensorflow/lite/python/analyzer.py @@ -96,11 +96,22 @@ def analyze(model_path=None, input_is_filepath = False if kwargs.get("experimental_use_mlir", False): - print( - wrap_converter.wrapped_flat_buffer_file_to_mlir( - tflite_model, input_is_filepath - ) - ) + try: + mlir_text = wrap_converter.wrapped_flat_buffer_file_to_mlir( + tflite_model, + input_is_filepath, + cl_options=[ + "-mlir-print-local-scope", + "-mlir-elide-elementsattrs-if-larger=16", + ], + ) + except TypeError: + # Fallback for environments where wrapped_flat_buffer_file_to_mlir does + # not yet accept cl_options (e.g. OSS LiteRT using external TF). + mlir_text = wrap_converter.wrapped_flat_buffer_file_to_mlir( + tflite_model, input_is_filepath + ) + print(mlir_text) else: print( _analyzer_wrapper.ModelAnalyzer(tflite_model, input_is_filepath, diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 9d319e54b894b4..4bee63b60d07c0 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -198,6 +198,67 @@ def get_options(): return [str(option) for option in list(OpsSet)] +@convert_phase(Component.CONVERT_TF_TO_TFLITE_MODEL, SubComponent.UNSPECIFIED) +def flatbuffer_to_mlir( + model_content, + input_is_filepath=False, + bytecode=False, + cl_options=None, +): + """Converts a TFLite FlatBuffer model to MLIR string or bytecode. + + Args: + model_content: A TFLite FlatBuffer as bytes, or a path to a TFLite file if + input_is_filepath is True. + input_is_filepath: If True, model_content is treated as a file path. + bytecode: If True, returns MLIR bytecode (.mlirc / .mlirbc) as bytes. + Otherwise returns textual MLIR as a string. + cl_options: Sequence of MLIR command-line printing options to forward. + + Returns: + str if bytecode=False, bytes if bytecode=True. + """ + return wrap_converter.wrapped_flat_buffer_file_to_mlir( + model_content, + input_is_filepath=input_is_filepath, + bytecode=bytecode, + cl_options=cl_options, + ) + + +@convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.UNSPECIFIED) +def mlir_to_flatbuffer( + mlir_content, + input_is_filepath=False, + emit_builtin_tflite_ops=True, + emit_select_tf_ops=False, + emit_custom_ops=True, + emit_stablehlo_ops=False, +): + """Converts an MLIR source (textual or bytecode) to a TFLite FlatBuffer. + + Args: + mlir_content: MLIR source as a string (textual IR) or bytes (bytecode), or a + file path if input_is_filepath is True. + input_is_filepath: If True, mlir_content is treated as a file path. + emit_builtin_tflite_ops: Whether to emit builtin TFLite operations. + emit_select_tf_ops: Whether to emit Select TF operations (Flex ops). + emit_custom_ops: Whether to allow custom operations. + emit_stablehlo_ops: Whether to serialize StableHLO operations. + + Returns: + TFLite FlatBuffer as bytes. + """ + return wrap_converter.wrapped_mlir_to_flat_buffer( + mlir_content, + input_is_filepath=input_is_filepath, + emit_builtin_tflite_ops=emit_builtin_tflite_ops, + emit_select_tf_ops=emit_select_tf_ops, + emit_custom_ops=emit_custom_ops, + emit_stablehlo_ops=emit_stablehlo_ops, + ) + + @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.QUANTIZE) def mlir_quantize( input_data_str, diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 0c6dfa0c9adbb9..65a6416d1324de 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -46,8 +46,10 @@ from tensorflow.lite.python.convert import convert_saved_model as _convert_saved_model from tensorflow.lite.python.convert import ConverterError # pylint: disable=unused-import from tensorflow.lite.python.convert import deduplicate_readonly_buffers as _deduplicate_readonly_buffers +from tensorflow.lite.python.convert import flatbuffer_to_mlir as _flatbuffer_to_mlir from tensorflow.lite.python.convert import mlir_quantize as _mlir_quantize from tensorflow.lite.python.convert import mlir_sparsify as _mlir_sparsify +from tensorflow.lite.python.convert import mlir_to_flatbuffer as _mlir_to_flatbuffer from tensorflow.lite.python.convert import OpsSet from tensorflow.lite.python.convert import toco_convert # pylint: disable=unused-import from tensorflow.lite.python.convert_phase import Component @@ -3543,3 +3545,4 @@ def from_keras_model_file( return TFLiteConverter.from_keras_model_file( model_file, input_arrays, input_shapes, output_arrays ) + diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index da0c23f3ac5937..bc3e0440313956 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -2904,5 +2904,43 @@ def testDeprecatedOptionWarning(self, optimization): logging.root.removeHandler(handler) +class FlatbufferMlirTranslationTest(TestModels): + + def testFlatbufferToMlirAndBack(self): + saved_model_dir = os.path.join( + self.get_temp_dir(), 'simple_savedmodel_mlir' + ) + with ops.Graph().as_default(): + with session.Session() as sess: + in_tensor = array_ops.placeholder( + shape=[1, 16], dtype=dtypes.float32, name='input' + ) + out_tensor = in_tensor + 1.0 + inputs = {'x': in_tensor} + outputs = {'z': out_tensor} + saved_model.simple_save(sess, saved_model_dir, inputs, outputs) + converter = lite.TFLiteConverter.from_saved_model(saved_model_dir) + tflite_model = converter.convert() + self.assertIsNotNone(tflite_model) + + mlir_str = lite._flatbuffer_to_mlir(tflite_model) + self.assertIsInstance(mlir_str, str) + self.assertIn('module', mlir_str) + + mlir_elided = lite._flatbuffer_to_mlir( + tflite_model, cl_options=['-mlir-elide-elementsattrs-if-larger=8'] + ) + self.assertIsInstance(mlir_elided, str) + + mlir_bc = lite._flatbuffer_to_mlir(tflite_model, bytecode=True) + self.assertIsInstance(mlir_bc, bytes) + + reloaded_tflite = lite._mlir_to_flatbuffer(mlir_str) + self.assertIsNotNone(reloaded_tflite) + + reloaded_tflite_from_bc = lite._mlir_to_flatbuffer(mlir_bc) + self.assertIsNotNone(reloaded_tflite_from_bc) + + if __name__ == '__main__': test.main() From 960a7fa5f8193573c13220fa61c12e6434a3fd6f Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 4 Sep 2026 21:59:13 -0700 Subject: [PATCH 11/12] Automated Code Change PiperOrigin-RevId: 976640785 --- third_party/xla/xla/python/ifrt/sharding.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/sharding.cc b/third_party/xla/xla/python/ifrt/sharding.cc index a641d51f29e315..eda18e355a077b 100644 --- a/third_party/xla/xla/python/ifrt/sharding.cc +++ b/third_party/xla/xla/python/ifrt/sharding.cc @@ -22,11 +22,9 @@ limitations under the License. #include #include #include -#include #include #include "absl/algorithm/container.h" -#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/status/status.h" From 01990a8ef1f82e62ffffc77ac0d1adc78dc5ee69 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 4 Sep 2026 23:42:10 -0700 Subject: [PATCH 12/12] Automated Code Change PiperOrigin-RevId: 976663042 --- .../python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc | 1 - third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc | 1 - third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc | 1 - third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc | 1 - .../python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc | 1 - third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc | 1 - third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc | 1 - 7 files changed, 7 deletions(-) diff --git a/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc index cb6ecf2d0d1857..139bba53c524ce 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/python/pjrt_ifrt/gpu_xla_executable_abi_version_serdes.h" #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc index 298cd13872229a..b2e91574307647 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/python/pjrt_ifrt/pjrt_client.h" #include -#include #include #include #include diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc index b2502d0330e6da..3f3d42e2edbe33 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_layout_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc index 74523156193a9f..6d0182b3ac60ea 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/reshard_impl_test_lib.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include #include -#include #include #include #include diff --git a/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc index e20a8d36e5c62e..68e5913ac29bbb 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/tpu_xla_executable_abi_version_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc index 373f8508abcb46..6ffda12e9c988b 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc index 9f6593a954619f..02bc31872caf94 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_serdes.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "absl/status/status.h"