Lowering Aten op to composite op instead of small ops #8502

Zantares · 2024-12-18T10:14:00Z

This PR is to solve the 2nd question in this issue: supports composite op in training.

Motivation

Composite op is beneficial for performance optimization and we aim to apply it to training too. . According to the response in the issue, the community has no plan to extend this to training currently... Thus, I created this draft PR to demonstrate our intention.

Detail

This PR alters the Aten op lowering process when there isn't a 1:1 mapping to XLA op. It uses composite call instead of small XLA ops. Later, in the optimization process, the composite call can be easily replaced with a custom kernel or decomposed.

This is still a draft PR and only Gelu is implemented as an example. If it gets accepted, here are some further suggestions:

Keep both the decomposed ops and the composite call implementation. Use a new env setting (e.g. XLA_COMPOSITE_OP) to enable this feature. Also, add an op list setting to define which ops can be composed.
Only retain the composite call implementation as it can be easily decomposed by StableHLO pass. User can control the behavior by turning the decompose pass on or off.

Example

import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
gelu = torch.nn.GELU(approximate="none")

x = torch.tensor([2.0], requires_grad=True, device=device)
y = gelu(x ** 2)
y.backward()

print(x.grad)

With this PR, the generated StableHLO is:

module @SyncTensorsGraph.43 attributes {mhlo.cross_program_prefetches = [], mhlo.input_output_alias = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func private @composite.gelu_backward.14(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
    %cst = stablehlo.constant dense<0.398942292> : tensor<1xf32>
    %cst_0 = stablehlo.constant dense<-5.000000e-01> : tensor<1xf32>
    %cst_1 = stablehlo.constant dense<5.000000e-01> : tensor<1xf32>
    %cst_2 = stablehlo.constant dense<1.000000e 00> : tensor<1xf32>
    %cst_3 = stablehlo.constant dense<0.707106769> : tensor<1xf32>
    %0 = stablehlo.multiply %arg1, %cst_3 : tensor<1xf32>
    %1 = stablehlo.custom_call @mhlo.erf(%0) {mhlo.attributes = {}, mhlo.version = 1 : i64} : (tensor<1xf32>) -> tensor<1xf32>
    %2 = stablehlo.add %1, %cst_2 : tensor<1xf32>
    %3 = stablehlo.multiply %2, %cst_1 : tensor<1xf32>
    %4 = stablehlo.multiply %arg1, %arg1 : tensor<1xf32>
    %5 = stablehlo.multiply %4, %cst_0 : tensor<1xf32>
    %6 = stablehlo.exponential %5 : tensor<1xf32>
    %7 = stablehlo.multiply %arg1, %6 : tensor<1xf32>
    %8 = stablehlo.multiply %7, %cst : tensor<1xf32>
    %9 = stablehlo.add %3, %8 : tensor<1xf32>
     = stablehlo.multiply %arg0, %9 : tensor<1xf32>
    return  : tensor<1xf32>
  }
  func.func @main(%arg0: tensor<f32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
    %cst = stablehlo.constant dense<1.000000e 00> : tensor<1xf32>
    %cst_0 = stablehlo.constant dense<2.000000e 00> : tensor<1xf32>
    %0 = stablehlo.power %arg1, %cst_0 : tensor<1xf32>
    %1 = stablehlo.composite "composite.gelu_backward" %cst, %0 {composite_attributes = {approximate = "none"}, decomposition = @composite.gelu_backward.14, version = 1 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
    %2 = stablehlo.power %arg1, %cst : tensor<1xf32>
    %3 = stablehlo.reshape %arg0 : (tensor<f32>) -> tensor<1xf32>
    %4 = stablehlo.multiply %2, %3 : tensor<1xf32>
    %5 = stablehlo.multiply %1, %4 : tensor<1xf32>
    return %5 : tensor<1xf32>
  }
}

lsy323 · 2024-12-20T04:55:56Z

torch_xla/csrc/ops/ops.cpp

+
+    // Building composite computation.
+    const std::string name = "composite.gelu";
+    const std::string attr = "{approximate = \"none\"}";


This is a dummy str for testing purpose?

This is a real op attribution for GELU: https://pytorch.org/docs/stable/generated/torch.nn.GELU.html#torch.nn.GELU

The available value of approximate is none or tanh. The lowering process checks this attribution and decides the sub lower function here. As my changes are in the sub lower function, I manually set this attribution.

It's a common process for composite op which has attributions (defined as non-tensor inputs for composite op, e.g. dim for Softmax).

Maybe I can get the attribution from XlaOp instead of manually setting strings, I will try.

What information is important for optimizations on this op? One option could be to have the composites implied value be none, and not generate a composite if tanh is used, of if the value is needed, then this looks good as-is, not a great API for making these in XLA, since they're an MLIR-first feature currently. If an MLIR dep is allowed in this file, you could build an MLIR dict and then dump to string before calling the XLA builder method.

lsy323 · 2024-12-20T04:56:18Z

torch_xla/csrc/ops/ops.cpp

+    // Building call to computation.
+    std::vector<xla::XlaOp> inputs{xla_input};
+    xla::XlaOp output = xla::CompositeCall(loctx->builder(), computation, inputs, name,
+                                           attr, /*version=*/1);


ditto for version

Yes this is for testing, I learned this setting from this XLA UT. I can remove it if it makes no scense.

lsy323 · 2024-12-20T04:57:05Z

Hi @Zantares, thanks for the PR! As long as the normal compilation/execution flow is not affected, I think this actually makes the HLO/StableHLO cleaner.

Not sure if you have access to TPU to see if the code example in the PR description would run on TPU as well. I applied TPU CI tag, could you please push again to see if TPU CI pass?

lsy323 · 2024-12-20T04:57:53Z

cc @GleasonK in case you know if the composite HLO op affects complication flow

Zantares · 2024-12-20T06:26:47Z

Hi @Zantares, thanks for the PR! As long as the normal compilation/execution flow is not affected, I think this actually makes the HLO/StableHLO cleaner.

Not sure if you have access to TPU to see if the code example in the PR description would run on TPU as well. I applied TPU CI tag, could you please push again to see if TPU CI pass?

Thanks! I will fix the format error and push it again to trigger TPU CI.

GleasonK · 2025-01-07T18:35:32Z

torch_xla/csrc/ops/ops.cpp

-    return node.ReturnOp(BuildGelu(xla_input), loctx);
+
+    // Building composite computation.
+    const std::string name = "composite.gelu";


The purpose of namespacing is to be able to tell who the maintainer or origin of a given composite is. If something changes about GELU (new attr, etc) who is on the hook to maintain it (for composites this is usually intended to be a vendor who has a library nvidia.some_op, aws.some_op, litert.some_op etc). The name composite doesn't answer this question.

For this I'd recommend ptxla.gelu or aten.gelu as names.

As an aside, I'd be careful representing converting every (or even many) aten operations using composites. They have a maintenance overhead in terms of needing to consider forward/backward compatibilty (i.e. the above thing about "what if gelu changes, who fixes?"), for some ops like gelu/softmax its probably ok, they don't tend to change much and usually look somewhat uniform.

For other aten ops that have very specific HW support, I'd recommend an approach that decentralizes composite ownership/maintenance, i.e. FX graph rewrite-as-composite API or make composite builder work for these use cases. This is how Google AI Edge uses composites today, they own the library and the compatibility for the (very small) subset of ATen ops that they have HW support for.

(cc @lsy323)

Thanks for the suggestions, here are some of my thoughts:

The namespace ptxla sounds quite good, and I will address it.

Regarding the compatibility/ownership of these ops in this PR, I believe it still belongs totorch-xla. The reason is that I didn't re-implement these ops from scratch, I simply wrapped the original implementation (such as BuildGelu and LowerSoftmax) with a composite call. W/O this PR, those implementations would need to be fixed if there were any changes in the semantics. At the current stage, I don't have any plans to introduce new composite ops that don't have an original implementation in torch-xla

Since this PR is aimed at resolving training issues, I'm uncertain whether the composite builder will work or not. Judging from the discussions in the attached issues, it appears that it might not work for training purposes. Could you please share some examples or guidance on:

FX graph rewrite-as-composite API or make composite builder work for these use cases.

Lowering Aten op to composite op instead of small ops.

49de8bc

Zantares mentioned this pull request Dec 18, 2024

2 questions for the composite op feature #8486

Open

lsy323 added the tpuci label Dec 20, 2024

lsy323 reviewed Dec 20, 2024

View reviewed changes

Zantares added 2 commits December 20, 2024 15:00

Fix format and remove version test info.

c9ee78e

Add Softmax and it's backward.

21501b7

Zantares changed the title ~~Draft: Lowering Aten op to composite op instead of small ops~~ Lowering Aten op to composite op instead of small ops Jan 2, 2025

GleasonK reviewed Jan 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lowering Aten op to composite op instead of small ops #8502

Lowering Aten op to composite op instead of small ops #8502

Zantares commented Dec 18, 2024

lsy323 Dec 20, 2024

Zantares Dec 20, 2024

Zantares Dec 20, 2024

GleasonK Jan 7, 2025 •

edited

Loading

lsy323 Dec 20, 2024

Zantares Dec 20, 2024

lsy323 commented Dec 20, 2024

lsy323 commented Dec 20, 2024

Zantares commented Dec 20, 2024

GleasonK Jan 7, 2025

GleasonK Jan 7, 2025 •

edited

Loading

GleasonK Jan 7, 2025

Zantares Jan 8, 2025

Lowering Aten op to composite op instead of small ops #8502

Are you sure you want to change the base?

Lowering Aten op to composite op instead of small ops #8502

Conversation

Zantares commented Dec 18, 2024

Motivation

Detail

Example

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GleasonK Jan 7, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lsy323 commented Dec 20, 2024

lsy323 commented Dec 20, 2024

Zantares commented Dec 20, 2024

Choose a reason for hiding this comment

GleasonK Jan 7, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GleasonK Jan 7, 2025 •

edited

Loading

GleasonK Jan 7, 2025 •

edited

Loading