[mhlo] AllGather variadic operands support #8151

apivovarov · 2024-01-03T23:35:34Z

Currently AllGather in HLO supports multiple operands/results, while MHLO only supports a single operand/result. This change addresses the parity gap by adding MHLO AllGather variadic operands support.

This change was inspired by previous commit 2457fc1 - [mhlo] AllReduce tuple support. Jun 7, 2023 by @GleasonK

Related commit:

PR-5740 [hlo] Add tuple input support to all-gather and reduce-scatter (Oct 16, 2023 by @jeffhataws)

@GleasonK @cheshire @burmako @jurahul @thomasjoerg Could you review this PR?

tdanyluk

It looks good based on the similarity to [mhlo] AllReduce tuple support.

I wouldn't mind if GleasonK also took a look.

jurahul · 2024-01-04T17:16:21Z

xla/mlir_hlo/mhlo/IR/hlo_ops.td

-def MHLO_AllGatherOp : MHLO_Op<"all_gather", [SameOperandsAndResultElementType]> {
+def MHLO_AllGatherOp : MHLO_Op<"all_gather", [
+    SameOperandsAndResultElementType,
+    SingleBlockImplicitTerminator<"ReturnOp">


The SingleBlockImplicitTerminator trait is not needed here. Unlike all-reduce op which has an attached region for the reduction computation, AllGather does not, so this trait is not relevant.

apivovarov · 2024-01-09T20:06:39Z

Hi Tamás, hi Rahul, seems that Kevin does not have time/resources to look at this PR. Can we merge it without his review?
@tdanyluk @jurahul

GleasonK

LGTM! Apologies for the delay, this fell off my radar with holidays!

GleasonK · 2024-01-09T20:13:19Z

xla/mlir_hlo/mhlo/IR/hlo_ops.cc

-                                getUseGlobalDeviceIds(), getResult());
+  if (getOperands().empty())
+    return emitOptionalError(getLoc(),
+                             "AllGather must have have at least one operand");


Do we verify anywhere that there are same number of operand and results?

I'm not sure if there's an upstream MLIR trait for this, if not then let's add another check here

Good catch! Added validation for that case to AllGatherOp::verify

if (getNumOperands() != getNumResults()) return emitOptionalError( getLoc(), "AllGather requires the same number of operands and results");

I am surprised there isn't an MLIR core trait to check this. If not, may be adding one for MLIR might be useful (not as a part of this PR though).

The list of OpTraits is here in LLVM mlir/include/mlir/IR/OpDefinition.h. I do not see a trait which compares Number of Operands and Results.
BTW, getOperands().empty() check might be redundant because SameOperandsAndResultElementType::verifyTrait calls OpTrait::impl::verifySameOperandsAndResultElementType(Operation *op) Operation.cpp which has the following

if (failed(verifyAtLeastNOperands(op, 1)) || failed(verifyAtLeastNResults(op, 1))) return failure();

GleasonK · 2024-01-09T20:19:42Z

xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc

-        all_gather_op.getOperand().getType().cast<TensorType>();
-    TensorType result_type = all_gather_op.getType();
+        all_gather_op.getOperand(0).getType().cast<TensorType>();
+    TensorType result_type = all_gather_op.getType(0).cast<TensorType>();
    if (!operand_type.hasStaticShape() || !result_type.hasStaticShape())
      return failure();
    if (operands.size() != 1) return failure();


Let's raise this check before the call to getOperand(0)

mhlo/IR/hlo_ops.td defines that MHLO_AllGatherOp has Variadic<MHLO_Tensor>:$operands,
I tried to use mhlo.all_gather() without operands in export.mlir - got error

error: unexpected error: 'mhlo.all_gather' op expected 1 or more operands, but found 0 %0:2 = "mhlo.all_gather"() {

mhlo parser validates that mhlo.all_gather() should have at least one operand.
It should be safe to use all_gather_op.getOperand(0) in mlir_hlo_to_hlo.cc::ExportXlaOp since mlir validation happens before ExportXlaOp call.

apivovarov · 2024-01-11T23:25:01Z

Just want to check if any other changes are needed for this PR? @GleasonK @jurahul @tdanyluk

apivovarov · 2024-01-16T20:18:06Z

This PR summary:

two official approvals from @GleasonK and @jurahul
"looks good" from @tdanyluk

Should we finally merge this PR? @ezhulenev @GleasonK @jurahul @tdanyluk

jurahul

SGTM.

GleasonK · 2024-01-16T21:04:40Z

xla/mlir_hlo/mhlo/IR/hlo_ops.cc

+        getLoc(),
+        "AllGather requires the same number of operands and results");
+
+  for (auto i = 0; i < getNumOperands(); ++i) {


mhlo/IR/hlo_ops.cc:2027: error: comparison of integers of different signs: 'int' and 'unsigned int' [-Werror,-Wsign-compare] for (auto i = 0; i < getNumOperands(); ++i) { ~ ^ ~~~~~~~~~~~~~~~~

We have an internal build failure on this. Let's use unsigned i to make sure the proper type is inferred.

fixed in 2 places

import/copybara still failing

I think the next action is on my end, need to update some internal callsites now that getOperand()/getResult() methods don't exist anymore (since they assume 1op/res). Will report back shortly.

Edit: Need to address the bazelrc feedback then I should be able to take care of the remaining integ work.

I removed .tf_configure.bazelrc and tools/python_bin_path.sh.
import/copybara check is OK now!

I believe I've made all the internal changes needed. Hopefully should have this merged tomorrow. I'll need to submit internally, once that happens this PR will be marked as closed and a GH commit attributed to you will be added to the main branch. Will give an update tomorrow!

Done: tensorflow/mlir-hlo@99207fc

Thank You, Kevin, for your help!

ddunl · 2024-01-16T22:50:12Z

.tf_configure.bazelrc

@@ -0,0 +1,16 @@
+build --action_env PYTHON_BIN_PATH="/usr/local/bin/python3.9"


I think this and tools/python_bin_path.sh should be removed. (I'll add .tf_configure.bazelrc to the gitignore, I'm surprised this hasn't happened before)

removed both. Sorry for the mess

No worries at all!!

@GleasonK

Imported from GitHub PR openxla/xla#8151 Currently AllGather in HLO supports multiple operands/results, while MHLO only supports a single operand/result. This change addresses the parity gap by adding MHLO AllGather variadic operands support. This change was inspired by previous commit [2457fc1](openxla/xla@2457fc1) - [mhlo] AllReduce tuple support. Jun 7, 2023 by @GleasonK Related commit: - [PR-5740](openxla/xla#5740) [hlo] Add tuple input support to all-gather and reduce-scatter (Oct 16, 2023 by @jeffhataws) @GleasonK @cheshire @burmako @jurahul @thomasjoerg Could you review this PR? Copybara import of the project: -- fb53ead74cbb40177a3680c8f807149d39c396b7 by Alexander Pivovarov <[email protected]>: [mhlo] AllGather variadic operands support Merging this change closes #8151 PiperOrigin-RevId: 599175008

@GleasonK

Imported from GitHub PR openxla/xla#8151 Currently AllGather in HLO supports multiple operands/results, while MHLO only supports a single operand/result. This change addresses the parity gap by adding MHLO AllGather variadic operands support. This change was inspired by previous commit [2457fc1](openxla/xla@2457fc1) - [mhlo] AllReduce tuple support. Jun 7, 2023 by @GleasonK Related commit: - [PR-5740](openxla/xla#5740) [hlo] Add tuple input support to all-gather and reduce-scatter (Oct 16, 2023 by @jeffhataws) @GleasonK @cheshire @burmako @jurahul @thomasjoerg Could you review this PR? Copybara import of the project: -- fb53ead74cbb40177a3680c8f807149d39c396b7 by Alexander Pivovarov <[email protected]>: [mhlo] AllGather variadic operands support Merging this change closes #8151 PiperOrigin-RevId: 599175008

apivovarov · 2024-01-17T23:30:17Z

Do I also need to update the documentation https://www.tensorflow.org/mlir/hlo_ops#mhloall_gather_mhloallgatherop?
What repo has the source of this page?

apivovarov · 2024-01-18T06:34:40Z

Seems that the documentation was updated automatically

operands | variadic of tensor of f8E4M3B11FNUZ type ...

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 3, 2024

github-actions bot assigned kamaljeeti and xla-rotation Jan 3, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 3, 2024

apivovarov force-pushed the mhlo_allgather_variadic branch from 2e3f49f to ed93860 Compare January 3, 2024 23:38

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 3, 2024

apivovarov changed the title ~~MHLO AllReduce variadic operands support~~ MHLO AllGather variadic operands support Jan 3, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 3, 2024

apivovarov changed the title ~~MHLO AllGather variadic operands support~~ [mhlo] AllGather variadic operands support Jan 3, 2024

kamaljeeti requested a review from tdanyluk January 4, 2024 04:10

tdanyluk requested a review from GleasonK January 4, 2024 15:32

tdanyluk reviewed Jan 4, 2024

View reviewed changes

jurahul reviewed Jan 4, 2024

View reviewed changes

apivovarov force-pushed the mhlo_allgather_variadic branch from ed93860 to 567b97b Compare January 4, 2024 19:30

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 4, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 4, 2024

GleasonK approved these changes Jan 9, 2024

View reviewed changes

apivovarov force-pushed the mhlo_allgather_variadic branch from 567b97b to 713913e Compare January 10, 2024 01:11

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 10, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 10, 2024

jurahul approved these changes Jan 16, 2024

View reviewed changes

apivovarov force-pushed the mhlo_allgather_variadic branch from 713913e to c9331d0 Compare January 16, 2024 20:24

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 16, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 16, 2024

GleasonK reviewed Jan 16, 2024

View reviewed changes

apivovarov force-pushed the mhlo_allgather_variadic branch from c9331d0 to 93e11aa Compare January 16, 2024 21:10

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 16, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 16, 2024

apivovarov force-pushed the mhlo_allgather_variadic branch from 93e11aa to e8bcccb Compare January 16, 2024 21:24

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 16, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 16, 2024

ddunl reviewed Jan 16, 2024

View reviewed changes

[mhlo] AllGather variadic operands support

fb53ead

apivovarov force-pushed the mhlo_allgather_variadic branch from e8bcccb to fb53ead Compare January 16, 2024 23:27

github-actions bot added the kokoro:force-run Forces CI to rerun label Jan 16, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Jan 16, 2024

GleasonK approved these changes Jan 16, 2024

View reviewed changes

ddunl approved these changes Jan 16, 2024

View reviewed changes

GleasonK self-assigned this Jan 17, 2024

copybara-service bot closed this in f6a6f97 Jan 17, 2024

apivovarov deleted the mhlo_allgather_variadic branch January 17, 2024 17:12

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mhlo] AllGather variadic operands support #8151

[mhlo] AllGather variadic operands support #8151

apivovarov commented Jan 3, 2024 •

edited

Loading

tdanyluk left a comment

jurahul Jan 4, 2024

apivovarov Jan 4, 2024

apivovarov commented Jan 9, 2024 •

edited

Loading

GleasonK left a comment

GleasonK Jan 9, 2024

apivovarov Jan 10, 2024

jurahul Jan 10, 2024

apivovarov Jan 10, 2024 •

edited

Loading

GleasonK Jan 9, 2024

apivovarov Jan 10, 2024 •

edited

Loading

apivovarov commented Jan 11, 2024

apivovarov commented Jan 16, 2024 •

edited

Loading

jurahul left a comment

GleasonK Jan 16, 2024

apivovarov Jan 16, 2024 •

edited

Loading

apivovarov Jan 16, 2024

GleasonK Jan 16, 2024 •

edited

Loading

apivovarov Jan 16, 2024 •

edited

Loading

GleasonK Jan 17, 2024

GleasonK Jan 17, 2024

apivovarov Jan 17, 2024

ddunl Jan 16, 2024

apivovarov Jan 16, 2024

ddunl Jan 16, 2024

apivovarov commented Jan 17, 2024

apivovarov commented Jan 18, 2024

		@@ -0,0 +1,16 @@
		build --action_env PYTHON_BIN_PATH="/usr/local/bin/python3.9"

[mhlo] AllGather variadic operands support #8151

[mhlo] AllGather variadic operands support #8151

Conversation

apivovarov commented Jan 3, 2024 • edited Loading

tdanyluk left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apivovarov commented Jan 9, 2024 • edited Loading

GleasonK left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apivovarov Jan 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apivovarov Jan 10, 2024 • edited Loading

Choose a reason for hiding this comment

apivovarov commented Jan 11, 2024

apivovarov commented Jan 16, 2024 • edited Loading

jurahul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apivovarov Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GleasonK Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

apivovarov Jan 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apivovarov commented Jan 17, 2024

apivovarov commented Jan 18, 2024

apivovarov commented Jan 3, 2024 •

edited

Loading

apivovarov commented Jan 9, 2024 •

edited

Loading

apivovarov Jan 10, 2024 •

edited

Loading

apivovarov Jan 10, 2024 •

edited

Loading

apivovarov commented Jan 16, 2024 •

edited

Loading

apivovarov Jan 16, 2024 •

edited

Loading

GleasonK Jan 16, 2024 •

edited

Loading

apivovarov Jan 16, 2024 •

edited

Loading