Semantics on channel_handle,replica_groups

[linuxlonelyeagle]

I'm new to stablehlo.I have some questions I would like to ask for help.You can read all_reduce.

• What is the semantics of channel_handle?
  I didn't really find much of a description of channel_handle too much.
• What is the semantics of replica_groups?
  My understanding is that processes are grouped, processes within the group communicate, and then produce results.
• What is the semantics of use_global_device_ids?
  I didn't really find much of a description of channel_handle too much.
• Test case on all_reduce
  Here are three test cases，@cross_replica, @cross_replica_and_partition, @flattened_ids,the problem is that the three test cases don't look different to me, and I don't understand the difference between the three examples even though I read the test cases.
  Finally, I hope someone can help me, I'm interested in stablehlo, thanks!

[sdasgup3]

You can find the description of channel handle, replica groups. https://github.com/openxla/stablehlo/blob/main/docs/spec.md#parallel-execution. Note that flatten_ids are enabled with use_global_device=true.

Test cases

It is the interoretation of the replica groups which are different in different tests cases. For the pretation you will find the descriptions handy.

[linuxlonelyeagle]

Thanks!

[linuxlonelyeagle]

I've read the parts you mentioned, but I seem to have a new problem now.replica_groups I can get through the properties of the communication operator. But the num_partitions value I don't know how I get it, which is very confusing to me.

[linuxlonelyeagle]

After my long quest, I'm making progress.In the following case, num_partitions is 2, but I'm having a bit of trouble understanding the following behavior.

module @cross_replica {
  func.func public @all_gather(%arg0 : tensor<2x2xi64>) -> tensor<2x4xi64> {
    %result = "stablehlo.all_gather"(%arg0) {
      all_gather_dim = 1 : i64,
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
    } : (tensor<2x2xi64>) -> tensor<2x4xi64>
    return %result : tensor<2x4xi64>
  }
  func.func public @main() {
    %0 = stablehlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi64>
    %1 = stablehlo.constant dense<[[5, 6], [7, 8]]> : tensor<2x2xi64>
    %results:4 = "interpreter.run_parallel"(%1, %1, %0, %1) {
      programs=[[@all_gather, @all_gather], [@all_gather,@all_gather]]
    } : (tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>)
    check.expect_eq_const %results#0, dense<[[5, 6, 1, 2],
                                             [7, 8, 3, 4]]> : tensor<2x4xi64>


    func.return
  }
}

[ghpvnist]

To get a better understanding, let's look at the full test output:

module @cross_replica {
  func.func public @all_gather(%arg0 : tensor<2x2xi64>) -> tensor<2x4xi64> {
    %result = "stablehlo.all_gather"(%arg0) {
      all_gather_dim = 1 : i64,
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
    } : (tensor<2x2xi64>) -> tensor<2x4xi64>
    return %result : tensor<2x4xi64>
  }
  func.func public @main() {
    %0 = stablehlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi64>
    %1 = stablehlo.constant dense<[[5, 6], [7, 8]]> : tensor<2x2xi64>
    %results:4 = "interpreter.run_parallel"(%1, %1, %0, %1) {
      programs=[[@all_gather, @all_gather], [@all_gather, @all_gather]]
    } : (tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>)
    check.expect_eq_const %results#0, dense<[[5, 6, 1, 2],
                                             [7, 8, 3, 4]]> : tensor<2x4xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 5, 6],
                                             [7, 8, 7, 8]]> : tensor<2x4xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 1, 2],
                                             [7, 8, 3, 4]]> : tensor<2x4xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 5, 6],
                                             [7, 8, 7, 8]]> : tensor<2x4xi64>
    func.return
  }
}

Firstly, note that the inputs to your program is %1, %1, %0, %1. For easier differentiation, I'll call the first input (%1) as input0, second input (%1) as input1, etc.

Note process_grid has the process_id in the form (replica_id, partition_id)
input0 is (0, 0) in the process_grid
input1 is (0, 1) in the process_grid
input2 is (1, 0) in the process_grid
input3 is (1, 1) in the process_grid

Using the cross_replica algorithm in the spec:
process_group = [[(0, 0), (1, 0)], [(0, 1), (1, 1)]]

Here, we see that the process_id of input0=(0, 0) and process_id of input2=(1, 0) are in one group, and they communicate with each other.
Similarly, process_id of input1=(0, 1) and process_id of input3=(1, 1) are in another group.

input0 and input2 concat (in that order) at dim=1 to get %results#0 and %results#2 as the same result is shared back to the process_id (0, 0) and (1, 0).
input1 and input3 concat (in that order) at dim=1 to get %results#1 and %results#3 as the same result is shared back to the process_id (0, 1) and (1, 1).

I hope this makes sense!

[linuxlonelyeagle]

I think I've figured it out.

[linuxlonelyeagle]

To get a better understanding, let's look at the full test output:

module @cross_replica {
  func.func public @all_gather(%arg0 : tensor<2x2xi64>) -> tensor<2x4xi64> {
    %result = "stablehlo.all_gather"(%arg0) {
      all_gather_dim = 1 : i64,
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
    } : (tensor<2x2xi64>) -> tensor<2x4xi64>
    return %result : tensor<2x4xi64>
  }
  func.func public @main() {
    %0 = stablehlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi64>
    %1 = stablehlo.constant dense<[[5, 6], [7, 8]]> : tensor<2x2xi64>
    %results:4 = "interpreter.run_parallel"(%1, %1, %0, %1) {
      programs=[[@all_gather, @all_gather], [@all_gather, @all_gather]]
    } : (tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>, tensor<2x4xi64>)
    check.expect_eq_const %results#0, dense<[[5, 6, 1, 2],
                                             [7, 8, 3, 4]]> : tensor<2x4xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 5, 6],
                                             [7, 8, 7, 8]]> : tensor<2x4xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 1, 2],
                                             [7, 8, 3, 4]]> : tensor<2x4xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 5, 6],
                                             [7, 8, 7, 8]]> : tensor<2x4xi64>
    func.return
  }
}

Firstly, note that the inputs to your program is %1, %1, %0, %1. For easier differentiation, I'll call the first input (%1) as input0, second input (%1) as input1, etc.

Note process_grid has the process_id in the form (replica_id, partition_id) input0 is (0, 0) in the process_grid input1 is (0, 1) in the process_grid input2 is (1, 0) in the process_grid input3 is (1, 1) in the process_grid

Using the cross_replica algorithm in the spec: process_group = [[(0, 0), (1, 0)], [(0, 1), (1, 1)]]

Here, we see that the process_id of input0=(0, 0) and process_id of input2=(1, 0) are in one group, and they communicate with each other. Similarly, process_id of input1=(0, 1) and process_id of input3=(1, 1) are in another group.

input0 and input2 concat (in that order) at dim=1 to get %results#0 and %results#2 as the same result is shared back to the process_id (0, 0) and (1, 0). input1 and input3 concat (in that order) at dim=1 to get %results#1 and %results#3 as the same result is shared back to the process_id (0, 1) and (1, 1).

I hope this makes sense!

With you what you said above, I wrote such an example and I have understood how to delineate which processes in a process group will communicate.

module @cross_replica {
  func.func public @all_gather(%arg0 : tensor<2x2xi64>) -> tensor<2x8xi64> {
    %result = "stablehlo.all_gather"(%arg0) {
      all_gather_dim = 1 : i64,
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>, 
      channel_handle = #stablehlo.channel_handle<handle=1, type=0>
    } : (tensor<2x2xi64>) -> tensor<2x8xi64>
    return %result : tensor<2x8xi64>
  }
  func.func public @main() {
    %0 = stablehlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi64>
    %1 = stablehlo.constant dense<[[5, 6], [7, 8]]> : tensor<2x2xi64>
    %results:4 = "interpreter.run_parallel"(%1, %1, %0, %1) {
      programs=[[@all_gather, @all_gather], [@all_gather,@all_gather]]
    } : (tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x8xi64>, tensor<2x8xi64>, tensor<2x8xi64>, tensor<2x8xi64>)
    check.expect_eq_const %results#0, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>

    func.return
  }
}

But I have my doubts about the results.Using the cross_replica_and_partition algorithm in the spec:
process_group = [[(0, 0), (1, 0), (0, 1), (1, 1)]]
input0 input2 input1 input3
input0 + input2 + input1 + input3=[[5, 6, 1, 2, 5, 6, 5, 6], [7, 8, 3, 4, 7, 8, 7, 8]
Maybe there's something I'm not understanding correctly.I really appreciate your help.

[ghpvnist]

To verify tests to be deterministic, the process ids are sorted before they are concatenated, and this is the behavior you're seeing right now. Revisiting the spec, it's somewhat implied that it should be concatenating in the order of the generated process groups instead. To that end, do you want to file a bug with the following information:

repro steps:

module @cross_replica {
  func.func public @all_gather(%arg0 : tensor<2x2xi64>) -> tensor<2x8xi64> {
    %result = "stablehlo.all_gather"(%arg0) {
      all_gather_dim = 1 : i64,
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>, 
      channel_handle = #stablehlo.channel_handle<handle=1, type=0>
    } : (tensor<2x2xi64>) -> tensor<2x8xi64>
    return %result : tensor<2x8xi64>
  }
  func.func public @main() {
    %0 = stablehlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi64>
    %1 = stablehlo.constant dense<[[5, 6], [7, 8]]> : tensor<2x2xi64>
    %results:4 = "interpreter.run_parallel"(%1, %1, %0, %1) {
      programs=[[@all_gather, @all_gather], [@all_gather,@all_gather]]
    } : (tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x8xi64>, tensor<2x8xi64>, tensor<2x8xi64>, tensor<2x8xi64>)
    check.expect_eq_const %results#0, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>

    func.return
  }
}

expected:

    check.expect_eq_const %results#0, dense<[[5, 6, 1, 2, 5, 6, 5, 6],
                                             [7, 8, 3, 4, 7, 8, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 1, 2, 5, 6, 5, 6],
                                             [7, 8, 3, 4, 7, 8, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 1, 2, 5, 6, 5, 6],
                                             [7, 8, 3, 4, 7, 8, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 1, 2, 5, 6, 5, 6],
                                             [7, 8, 3, 4, 7, 8, 7, 8]]> : tensor<2x8xi64>

result:

    check.expect_eq_const %results#0, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#1, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#2, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>
    check.expect_eq_const %results#3, dense<[[5, 6, 5, 6, 1, 2, 5, 6],
                                             [7, 8, 7, 8, 3, 4, 7, 8]]> : tensor<2x8xi64>

Thanks for bringing this up! I have a fix ready :)

[ghpvnist]

Thanks for filing the issue @linuxlonelyeagle, the fix is now merged in #1934!