Skip to content

Support multiple StableHLO Composite outputs#6295

Merged
chunnienc merged 6 commits intomasterfrom
chunnienc/hlfb-mutli-outputs
Feb 4, 2024
Merged

Support multiple StableHLO Composite outputs#6295
chunnienc merged 6 commits intomasterfrom
chunnienc/hlfb-mutli-outputs

Conversation

@chunnienc
Copy link
Copy Markdown
Collaborator

@chunnienc chunnienc commented Jan 11, 2024 •
edited
Loading

This PR adds the support for marking multiple tensors in StableHLOCompositeBuilder.mark_outputs.

Example:

class Model(torch.nn.Module):

      def __init__(self):
        super().__init__()

      def forward(self, x, y):
        builder = StableHLOCompositeBuilder("sample_composite")
        x, y = builder.mark_inputs(x, y)
        a = x + y
        b = x - y
        c = x + 1
        a, b, c = builder.mark_outputs(a, b, c)
        return a + b + c

SHLO:

module @IrToHlo.42 attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) -> tensor<5x5xf32> {
    %0:3 = stablehlo.custom_call @stablehlo.composite(%arg0, %arg1) {called_computations = [@sample_composite.impl], composite.backend_config = {attributes = {}, name = "sample_composite"}} : (tensor<5x5xf32>, tensor<5x5xf32>) -> (tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>)
    %1 = stablehlo.add %0#0, %0#1 : tensor<5x5xf32>
    %2 = stablehlo.add %1, %0#2 : tensor<5x5xf32>
    return %2 : tensor<5x5xf32>
  }
  func.func private @sample_composite.impl(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) -> (tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>) {
    %0 = stablehlo.constant dense<1.000000e+00> : tensor<5x5xf32>
    %1 = stablehlo.add %arg0, %arg1 : tensor<5x5xf32>
    %2 = stablehlo.subtract %arg0, %arg1 : tensor<5x5xf32>
    %3 = stablehlo.add %arg0, %0 : tensor<5x5xf32>
    return %1, %2, %3 : tensor<5x5xf32>, tensor<5x5xf32>, tensor<5x5xf32>
  }
}

@chunnienc chunnienc requested a review from lsy323 January 11, 2024 08:15
@miladm
Copy link
Copy Markdown
Collaborator

miladm commented Jan 11, 2024

Suggesting to include references to issues on public URLs on GH. Thanks very much.

@chunnienc chunnienc marked this pull request as ready for review February 1, 2024 06:17
@chunnienc chunnienc requested a review from qihqi February 1, 2024 18:09
lsy323
lsy323 approved these changes Feb 2, 2024
View reviewed changes
Copy link
Copy Markdown
Collaborator

@lsy323 lsy323 left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, left a few comments/questions. Thanks @chunnienc!

Comment thread torch_xla/csrc/runtime/stablehlo_composite_helper.cc
Comment thread torch_xla/csrc/runtime/stablehlo_composite_helper.cc
@chunnienc chunnienc merged commit 535d398 into master Feb 4, 2024
amithrm pushed a commit to amithrm/xla that referenced this pull request Mar 1, 2024
@chunnienc @amithrm
Support multiple StableHLO Composite outputs (pytorch#6295)
225b74c
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
@chunnienc @bhavya01
Support multiple StableHLO Composite outputs (#6295)
f5d6f69
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@lsy323 lsy323 lsy323 approved these changes

@qihqi qihqi Awaiting requested review from qihqi

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@chunnienc @miladm @lsy323