Make collective select folder convert-aware

PiperOrigin-RevId: 708378635

Author: frgossen

xla/service/gpu/transforms/BUILD

@@ -530,6 +530,7 @@ cc_library(
     hdrs = ["collective_select_folder.h"],
     deps = [
         "//xla:comparison_util",
+        "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/service:collective_ops_utils",

xla/service/gpu/transforms/collective_select_folder.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/collective_ops_utils.h"
+#include "xla/shape_util.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/statusor.h"
@@ -51,12 +52,20 @@ struct FoldableSelect {
   HloInstruction* false_operand;
 };
 
+const HloInstruction* FindInnerScalarOp(const HloInstruction* inst) {
+  while (inst->opcode() == HloOpcode::kConvert ||
+         inst->opcode() == HloOpcode::kBroadcast) {
+    inst = inst->operand(0);
+  }
+  return inst;
+}
+
 // Matches foldable select ops that we can analyse and returns handy references
 // to %constant, %true_operand, %false_operand of the op. Matches, e.g.,
 //
 // ```
 // select(
-//     broadcast(compare(partition-id(), constant)),
+//     broadcast(compare(convert(partition-id()), constant)),
 //     true_operand,
 //     false_operand)
 // ```
@@ -65,7 +74,7 @@ struct FoldableSelect {
 //
 // ```
 // select(
-//     compare(partition-id(), constant),
+//     compare(replica-id(), constant),
 //     true_operand,
 //     false_operand)
 // ```
@@ -74,21 +83,22 @@ std::optional<FoldableSelect> MatchFoldableSelect(HloInstruction* select) {
     return std::nullopt;
   }
 
-  // Match select predicate (may be broadcasted).
-  const HloInstruction* predicate_candidate = select->operand(0);
-  if (HloPredicateIsOp<HloOpcode::kBroadcast>(predicate_candidate))
-    predicate_candidate = predicate_candidate->operand(0);
+  // Match select predicate.
+  const HloInstruction* predicate_candidate =
+      FindInnerScalarOp(select->operand(0));
   const HloCompareInstruction* compare =
       DynCast<HloCompareInstruction>(predicate_candidate);
-  if (compare == nullptr) return std::nullopt;
+  if (compare == nullptr) {
+    return std::nullopt;
+  }
   if (compare->direction() != Comparison::Direction::kEq &&
       compare->direction() != Comparison::Direction::kNe) {
     return std::nullopt;
   }
 
   // Find replica-id or partition-id op and constant op, swap if needed.
-  const HloInstruction* id_op = compare->operand(0);
-  const HloInstruction* constant_op = compare->operand(1);
+  const HloInstruction* id_op = FindInnerScalarOp(compare->operand(0));
+  const HloInstruction* constant_op = FindInnerScalarOp(compare->operand(1));
   if (HloPredicateIsNotOp<HloOpcode::kConstant>(constant_op)) {
     std::swap(id_op, constant_op);
   }
@@ -104,10 +114,14 @@ std::optional<FoldableSelect> MatchFoldableSelect(HloInstruction* select) {
   }
 
   // Match constant.
-  if (HloPredicateIsNotOp<HloOpcode::kConstant>(constant_op))
+  if (HloPredicateIsNotOp<HloOpcode::kConstant>(constant_op) ||
+      !ShapeUtil::IsScalar(constant_op->shape())) {
     return std::nullopt;
+  }
   std::optional<int64_t> constant_id = constant_op->literal().GetFirstInteger();
-  if (!constant_id.has_value()) return std::nullopt;
+  if (!constant_id.has_value()) {
+    return std::nullopt;
+  }
   return FoldableSelect{compare->direction(), *constant_id, collective_mode,
                         select->mutable_operand(1), select->mutable_operand(2)};
 }

xla/service/gpu/transforms/collective_select_folder_test.cc

@@ -423,7 +423,7 @@ TEST_F(CollectiveSelectFolderTest,
     }
   )";
 
-  TF_ASSERT_OK_AND_ASSIGN(auto module,
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           RunAndCheckHloRewrite(kHlo, CollectiveSelectFolder(),
                                                 /*expect_change=*/true));
   const absl::string_view kExpected = R"(
@@ -449,5 +449,40 @@ TEST_F(CollectiveSelectFolderTest,
   EXPECT_TRUE(filecheck_result);
 }
 
+TEST_F(CollectiveSelectFolderTest, DtypeConvertedPartitionId) {
+  const absl::string_view kHlo = R"(
+    HloModule test
+
+    ENTRY computation {
+      param = (f32[1,1,28672,2048]{3,2,1,0}, f32[1,1,28672,2048]{3,2,1,0})
+          parameter(0)
+      get-tuple-element-a = f32[1,1,28672,2048]{3,2,1,0}
+          get-tuple-element(param), index=0
+      get-tuple-element-b = f32[1,1,28672,2048]{3,2,1,0}
+          get-tuple-element(param), index=1
+      partition-id.1 = u32[] partition-id()
+      convert = s32[] convert(partition-id.1)
+      constant.148 = s32[] constant(3)
+      compare.83 = pred[] compare(convert, constant.148), direction=EQ
+      select.33 = f32[1,1,28672,2048]{3,2,1,0} select(compare.83,
+          get-tuple-element-a, get-tuple-element-b)
+      ROOT cp-a = f32[1,1,28672,2048]{3,2,1,0} collective-permute(select.33),
+          channel_id=1, source_target_pairs={{3,0}}
+    }
+  )";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          RunAndCheckHloRewrite(kHlo, CollectiveSelectFolder(),
+                                                /*expect_change=*/true));
+  const absl::string_view kExpected = R"(
+    // CHECK: %[[PARAM:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[DATA_A:.*]] = {{.*}} get-tuple-element({{.*}} %[[PARAM]]), index=0
+    // CHECK: ROOT %[[DATA_A_:.*]] = {{.*}} collective-permute({{.*}} %[[DATA_A]])
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(bool filecheck_result,
+                          RunFileCheck(module->ToString(), kExpected));
+  EXPECT_TRUE(filecheck_result);
+}
+
 }  // namespace
 }  // namespace xla