openxla · burmako · Dec 20, 2022 · Dec 20, 2022 · Dec 20, 2022 · Dec 20, 2022
diff --git a/stablehlo/dialect/StablehloOps.cpp b/stablehlo/dialect/StablehloOps.cpp
@@ -1965,6 +1965,16 @@ LogicalResult AllToAllOp::inferReturnTypeComponents(
   return success();
 }
 
+void AllToAllOp::build(OpBuilder& odsBuilder, OperationState& odsState,
+                       Type resultType, Value operand,
+                       IntegerAttr splitDimension, IntegerAttr concatDimension,
+                       IntegerAttr splitCount,
+                       DenseIntElementsAttr replicaGroups) {
+  AllToAllOp::build(odsBuilder, odsState, resultType, operand, splitDimension,
+                    concatDimension, splitCount, replicaGroups,
+                    /*channel_handle=*/nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // AllGatherOp
 //===----------------------------------------------------------------------===//

diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td
@@ -1438,9 +1438,20 @@ def StableHLO_AllToAllOp : StableHLO_Op<"all_to_all",
     I64Attr:$split_dimension,
     I64Attr:$concat_dimension,
     I64Attr:$split_count,
-    I64ElementsAttr:$replica_groups
+    I64ElementsAttr:$replica_groups,
+    OptionalAttr<StableHLO_ChannelHandle>:$channel_handle
   );
   let results = (outs HLO_Tensor);
+
+  // channel_handle is only used for the SPMD partitioner, so we add a
+  // simplified builder method for convenience.
+  let builders = [
+    OpBuilder<(ins
+      "::mlir::Type":$result_type, "::mlir::Value":$operand,
+      "::mlir::IntegerAttr": $split_dimension,
+      "::mlir::IntegerAttr": $concat_dimension,
+      "::mlir::IntegerAttr": $split_count,
+      "::mlir::DenseIntElementsAttr": $replica_groups)>];
 }
 
 def StableHLO_ReduceOp: StableHLO_ShapedInterfaceOp<"reduce", [

diff --git a/stablehlo/dialect/VhloOps.td b/stablehlo/dialect/VhloOps.td
@@ -36,6 +36,7 @@ def VHLO_Dialect : Dialect {
              Add CollectivePermuteOp::channel_handle: https://github.com/openxla/stablehlo/pull/388.
              Add CustomCallOp::output_operand_aliases: https://github.com/openxla/stablehlo/pull/403.
              Add PartitionIdOp: https://github.com/openxla/stablehlo/pull/673.
+             Add AllToAllOp::channel_handle: TBD.
   }];
 
   let useDefaultAttributePrinterParser = 0;
@@ -401,7 +402,7 @@ def VHLO_ReduceScatterOpV1 : VHLO_Op<"reduce_scatter"> {
   let results = (outs VHLO_AnyType);
 }
 
-def VHLO_AllToAllOpV1 : VHLO_Op<"all_to_all"> {
+def VHLO_AllToAllOpV1 : VHLO_Op<"all_to_all", "0.3.0", "0.3.0"> {
   let arguments = (ins
     VHLO_AnyType:$operand,
     VHLO_AnyAttr:$split_dimension,
@@ -412,6 +413,18 @@ def VHLO_AllToAllOpV1 : VHLO_Op<"all_to_all"> {
   let results = (outs VHLO_AnyType);
 }
 
+def VHLO_AllToAllOpV2 : VHLO_Op<"all_to_all_v2", "0.4.0"> {
+  let arguments = (ins
+    VHLO_AnyType:$operand,
+    VHLO_AnyAttr:$split_dimension,
+    VHLO_AnyAttr:$concat_dimension,
+    VHLO_AnyAttr:$split_count,
+    VHLO_AnyAttr:$replica_groups,
+    OptionalAttr<VHLO_AnyAttr>:$channel_handle
+  );
+  let results = (outs VHLO_AnyType);
+}
+
 def VHLO_ReduceOpV1 : VHLO_Op<"reduce", "0.3.0", "current", [SameVariadicOperandSize]> {
   let arguments = (ins
     Variadic<VHLO_AnyType>:$inputs,

diff --git a/stablehlo/tests/legalize_stablehlo_to_vhlo.mlir b/stablehlo/tests/legalize_stablehlo_to_vhlo.mlir
@@ -380,7 +380,8 @@ func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
 // CHECK-LABEL: "op_all_reduce"
 
 func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-  //               CHECK: "vhlo.all_to_all"(%arg0) {
+  //               CHECK: "vhlo.all_to_all_v2"(%arg0) {
+  //          CHECK-SAME:   channel_handle = #vhlo.channel_handle<handle = 1, type = 0>
   //          CHECK-SAME:   concat_dimension = 0 : i64,
   // CHECK-SAME{LITERAL}:   replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
   //          CHECK-SAME:   split_count = 4 : i64,
@@ -390,7 +391,8 @@ func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
     split_dimension = 1 : i64,
     concat_dimension = 0 : i64,
     split_count = 4 : i64,
-    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
   } : (tensor<4x16xf32>) -> tensor<16x4xf32>
   func.return %0 : tensor<16x4xf32>
 }

diff --git a/stablehlo/tests/ops_stablehlo.mlir b/stablehlo/tests/ops_stablehlo.mlir
@@ -429,7 +429,8 @@ func.func @alltoall(%data: tensor<4x16xf32>) -> tensor<16x4xf32> {
     split_dimension = 1 : i64,
     concat_dimension = 0 : i64,
     split_count = 4 : i64,
-    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
   } : (tensor<4x16xf32>) -> tensor<16x4xf32>
   func.return %0 : tensor<16x4xf32>
 }

diff --git a/stablehlo/tests/vhlo_to_version_downgrade.mlir b/stablehlo/tests/vhlo_to_version_downgrade.mlir
@@ -1,5 +1,16 @@
 // RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=0.3.0' %s | FileCheck %s
 
+// CHECK-LABEL: @all_to_all_to_v1
+func.func @all_to_all_to_v1(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
+  // CHECK-NEXT: %0 = "vhlo.all_to_all"(%arg0)
+  %0 = "stablehlo.all_to_all"(%arg0) {
+    split_dimension = 1 : i64,
+    concat_dimension = 0 : i64,
+    split_count = 4 : i64,
+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
+  func.return %0 : tensor<16x4xf32>
+}
 
 // CHECK-LABEL: @all_gather_to_v1
 func.func @all_gather_to_v1(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {

diff --git a/stablehlo/tests/vhlo_to_version_downgrade_invalid.mlir b/stablehlo/tests/vhlo_to_version_downgrade_invalid.mlir
@@ -80,3 +80,18 @@ func.func @invalid_program_unknown_op(%arg0 : tensor<f32>) -> (tensor<f32>) {
   %0 = "vhlo.unknown_op"(%arg0) : (tensor<f32>) -> tensor<f32>
   func.return
 }
+
+// -----
+
+func.func @all_to_all_to_v1(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
+  // expected-error @+2 {{failed to downgrade vhlo.all_to_all_v2, op has a non-empty channel_handle attribute}}
+  // expected-error @+1 {{failed to legalize operation 'vhlo.all_to_all_v2' that was explicitly marked illegal}}
+  %0 = "stablehlo.all_to_all"(%arg0) {
+    split_dimension = 1 : i64,
+    concat_dimension = 0 : i64,
+    split_count = 4 : i64,
+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
+  func.return %0 : tensor<16x4xf32>
+}
diff --git a/stablehlo/tests/vhlo_to_version_upgrade.mlir b/stablehlo/tests/vhlo_to_version_upgrade.mlir
@@ -1,23 +1,44 @@
 // RUN: stablehlo-opt --vhlo-to-version='target=0.4.0' %s | FileCheck %s
 // RUN: stablehlo-opt --vhlo-to-version='target=current' %s | FileCheck %s
 
+// CHECK-LABEL: @all_to_all_to_v2
+func.func @all_to_all_to_v2(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
+  // CHECK-NEXT: %0 = "vhlo.all_to_all_v2"(%arg0)
+  %0 = "vhlo.all_to_all"(%arg0) {
+    split_dimension = 1 : i64,
+    concat_dimension = 0 : i64,
+    split_count = 4 : i64,
+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
+  func.return %0 : tensor<16x4xf32>
+}
+
 // CHECK-LABEL: @all_gather_to_v2
 func.func @all_gather_to_v2(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
   // CHECK-NEXT: %0 = "vhlo.all_gather_v2"(%arg0)
-  %0 = "vhlo.all_gather"(%arg0) {all_gather_dim = 1 : i64, channel_handle = #vhlo.channel_handle<handle = 0, type = 0>, replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>} : (tensor<16x8xf32>) -> tensor<16x16xf32>
+  %0 = "vhlo.all_gather"(%arg0) {
+    all_gather_dim = 1 : i64,
+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
+    channel_handle = #vhlo.channel_handle<handle = 0, type = 0>
+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
   return %0 : tensor<16x16xf32>
 }
 
 // CHECK-LABEL: @collective_permute_to_v2
 func.func @collective_permute_to_v2(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
   // CHECK-NEXT: %0 = "vhlo.collective_permute_v2"(%arg0)
-  %0 = "vhlo.collective_permute"(%arg0) {source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>} : (tensor<16x8xf32>) -> tensor<16x8xf32>
+  %0 = "vhlo.collective_permute"(%arg0) {
+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
   return %0 : tensor<16x8xf32>
 }
 
 // CHECK-LABEL: @custom_call_to_v2
 func.func @custom_call_to_v2(%arg0: tensor<2xi1>) -> tensor<2xi1> {
   // CHECK-NEXT: %0 = "vhlo.custom_call_v2"(%arg0)
-  %0 = "vhlo.custom_call"(%arg0) {backend_config = "", call_target_name = "foo"} : (tensor<2xi1>) -> tensor<2xi1>
+  %0 = "vhlo.custom_call"(%arg0) {
+    call_target_name = "foo",
+    backend_config = ""
+  } : (tensor<2xi1>) -> tensor<2xi1>
   return %0 : tensor<2xi1>
 }
diff --git a/stablehlo/transforms/MapStablehloToVhlo.h b/stablehlo/transforms/MapStablehloToVhlo.h
@@ -52,7 +52,7 @@ MAP_STABLEHLO_TO_VHLO(AddOp, V1)
 MAP_STABLEHLO_TO_VHLO(AfterAllOp, V1)
 MAP_STABLEHLO_TO_VHLO(AllGatherOp, V2)
 MAP_STABLEHLO_TO_VHLO(AllReduceOp, V1)
-MAP_STABLEHLO_TO_VHLO(AllToAllOp, V1)
+MAP_STABLEHLO_TO_VHLO(AllToAllOp, V2)
 MAP_STABLEHLO_TO_VHLO(AndOp, V1)
 MAP_STABLEHLO_TO_VHLO(Atan2Op, V1)
 MAP_STABLEHLO_TO_VHLO(BatchNormGradOp, V1)

diff --git a/stablehlo/transforms/VhloToVersion.cpp b/stablehlo/transforms/VhloToVersion.cpp
@@ -313,6 +313,28 @@ struct AllGatherOpV2ToV1
   }
 };
 
+// vhlo.all_to_all --> vhlo.all_to_all_v2
+struct AllToAllOpV1ToV2
+    : public VersionConversionPattern<AllToAllOpV1, AllToAllOpV2> {
+  using VersionConversionPattern::VersionConversionPattern;
+  LogicalResult prepareOpForConversion(AllToAllOpV1) const final {
+    return success();
+  }
+};
+
+// vhlo.all_to_all_v2 --> vhlo.all_to_all
+struct AllToAllOpV2ToV1
+    : public VersionConversionPattern<AllToAllOpV2, AllToAllOpV1> {
+  using VersionConversionPattern::VersionConversionPattern;
+  LogicalResult prepareOpForConversion(AllToAllOpV2 op) const final {
+    if (op.getChannelHandle().has_value()) {
+      return emitDowngradeError(op,
+                                "op has a non-empty channel_handle attribute");
+    }
+    return success();
+  }
+};
+
 }  // namespace
 }  // namespace vhlo
 
@@ -326,6 +348,8 @@ void populateVhloToVersionPatterns(RewritePatternSet* patterns,
   patterns->add<vhlo::CollectivePermuteOpV2ToV1>(*converter, context);
   patterns->add<vhlo::AllGatherOpV1ToV2>(*converter, context);
   patterns->add<vhlo::AllGatherOpV2ToV1>(*converter, context);
+  patterns->add<vhlo::AllToAllOpV1ToV2>(*converter, context);
+  patterns->add<vhlo::AllToAllOpV2ToV1>(*converter, context);
 }
 
 }  // namespace stablehlo