Repository: alibaba/BladeDISC (language: python)

Issue:
support inplace operator in TorchBlade
Pytorch inplace operators are widely used in model training/inference, just as KV cache in LLM, TorchBlade should support these
inplace operator.
 
Just as the following simple PyTorch program with in-place add operator `aten.add_`:
``` python
def func(add : Tensor, value : Tensor) -> Tensor:
  add.add_(value)
  return add
```
torchdynamo can trace the above program into value-semantic operators and insert `aten._copy` at the back:

``` python
add: f32[8, 64] = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg1_1 = None
# No stacktrace found for the following nodes
copy_: f32[8, 64] = torch.ops.aten.copy_.default(arg0_1, add);  arg0_1 = None
```

The naive operation in buffer level can be represent as:
``` mlir
lmhlo.add %arg0, %arg1, %1
lmhlo.copy %1, %arg0
```

but in actuality the `lmhlo.copy ` causes overhead because of the necessary d2d memcpy, we can remove this copy and store the result element on the input buffer:

``` mlir
lmhlo.add %arg0, %arg1, %arg0
``` 



Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/fusion.cpp b/pytorch_blade/pytorch_blade/compiler/jit/fusion.cpp
index f96fea44d9e..985fb977b3d 100644
--- a/pytorch_blade/pytorch_blade/compiler/jit/fusion.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/jit/fusion.cpp
@@ -140,16 +140,9 @@ void set_tensor_shape(
     torch::jit::Value* val,
     const std::vector<int64_t>& dims) {
   auto tensor_type = val->type()->cast<c10::TensorType>();
-  if (tensor_type) {
-    val->setType(c10::TensorType::create(
-        tensor_type->scalarType(),
-        tensor_type->device(),
-        c10::SymbolicShape(dims),
-        c10::VaryingShape<c10::Stride>(dims.size()),
-        tensor_type->requires_grad()));
-    return;
-  }
-  TORCH_CHECK(false, "input value should be tensor type");
+  if (!tensor_type)
+    TORCH_CHECK(false, "input value should be tensor type");
+  val->setType(tensor_type->withSizes(dims));
 }
 
 } // namespace blade
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
index 556bd283c1f..8331579c30d 100644
--- a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
@@ -1060,6 +1060,7 @@ class ShapePropagator : public PropertyPropBase {
             "aten::alias(Tensor self) -> Tensor",
             "aten::zero_(Tensor self) -> Tensor",
             "aten::tanh_backward(Tensor grad_output, Tensor output) -> Tensor",
+            "aten::copy_(Tensor(a!) self, Tensor src, bool non_blocking=False) -> Tensor(a!)",
 #ifdef TORCH_BLADE_BUILD_QUANTIZATION
             "torch_blade::fake_quant(Tensor _0, Tensor _1, Tensor _2, int _3, int _4, int _5, int[] _6, bool _7, bool _8, bool _9, bool _10) -> Tensor",
             "torch_blade::placeholder(Tensor _0) -> (Tensor _0)",
@@ -2725,6 +2726,8 @@ class ShapePropagator : public PropertyPropBase {
           int64_t end = endOptional.value() != c10::nullopt
               ? node->get<int64_t>(attr::end).value()
               : INT64_MAX;
+          start =
+              at::maybe_wrap_dim(start, new_sizes[dim].static_size(), false);
           int64_t step = node->get<int64_t>(attr::step).value();
           if (end >= new_sizes[dim].static_size())
             end = new_sizes[dim].static_size();
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
index cf2f83bc1f3..e647908c6cd 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
@@ -43,7 +43,6 @@ void mlir::torch::createDiscTorchBackendToMhloBackendPipeline(
   ::mlir::torch::Torch::TorchLoweringPipelineOptions funcOptions;
   funcOptions.decompose = false;
   ::mlir::torch::createDiscTorchFunctionToTorchBackendPipeline(pm, funcOptions);
-
   // Apply pdl pattern match
   std::string disc_torch_pdl_files;
   std::string disc_torch_pdll_include_dirs;
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
index b5ab1c0d3c1..dbcd1c7e0c7 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
@@ -765,7 +765,7 @@ LogicalResult ConvertAtenOp<AtenSizeIntOp>::matchAndRewrite(
     OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   // Not a tensor type.
-  auto selfType = adaptor.getSelf().getType().dyn_cast<TensorType>();
+  auto selfType = adaptor.getSelf().getType().cast<RankedTensorType>();
   if (!selfType)
     return op.emitError("Only tensor types are currently supported");
 
@@ -1465,6 +1465,33 @@ LogicalResult ConvertAtenOp<AtenConstantPadNdOp>::matchAndRewrite(
 }
 } // namespace
 
+namespace {
+template <typename T>
+Value backtraceOperand(Value operand) {
+  auto op = operand.getDefiningOp();
+  if (op && mlir::isa<T>(op)) {
+    return op->getOperand(0);
+  }
+  return operand;
+}
+
+template <>
+LogicalResult ConvertAtenOp<OverwriteTensorContentsOp>::matchAndRewrite(
+    OverwriteTensorContentsOp op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  auto loc = op.getLoc();
+  auto operands = op.getOperands();
+  auto value = operands[0];
+  auto overwriten = operands[1];
+  overwriten = backtraceOperand<CopyToNonValueTensorOp>(overwriten);
+  overwriten = rewriter.create<ToBuiltinTensorOp>(loc, overwriten);
+  value = rewriter.create<ToBuiltinTensorOp>(loc, value);
+  rewriter.replaceOpWithNewOp<mhlo_disc::ArgsMutationOp>(op, value, overwriten);
+  return success();
+}
+} //  namespace
+
 namespace {
 template <>
 LogicalResult ConvertAtenOp<AtenSqueezeDimOp>::matchAndRewrite(
@@ -1538,6 +1565,7 @@ class DiscConvertTorchToMhlo
     target.addLegalDialect<
         chlo::ChloDialect,
         mhlo::MhloDialect,
+        mhlo_disc::MhloDiscDialect,
         tensor::TensorDialect,
         arith::ArithDialect,
         Torch::TorchDialect>();
@@ -1631,6 +1659,7 @@ class DiscConvertTorchToMhlo
     INSERT_ATENOP_PATTERN(AtenSqueezeDimOp);
     INSERT_ATENOP_PATTERN(AtenNegIntOp);
     INSERT_ATENOP_PATTERN(AtenFillScalarOp);
+    INSERT_ATENOP_PATTERN(OverwriteTensorContentsOp);
 #undef INSERT_ATENOP_PATTERN
 
 #define INSERT_BINARY_BROADCAST_PATTERN(AtenOp, MhloOp)       \
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
index 24392141b03..178a5a3e312 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
@@ -54,6 +54,7 @@ class ApplyValueSemanticsPass
     : public TorchConversion::ApplyValueSemanticsBase<ApplyValueSemanticsPass> {
   void runOnOperation() override {
     MLIRContext* context = &getContext();
+    mlir::OpBuilder rewriter(context);
     auto func = getOperation();
 
     SmallVector<Operation*> deadOps;
@@ -75,8 +76,8 @@ class ApplyValueSemanticsPass
     for (auto op : deadOps) {
       op->erase();
     }
-
     reduceTensorConversions(func);
+
     deadOps.clear();
     func.walk([&](mlir::Operation* op) {
       if (isa<TensorStaticInfoCastOp>(op)) {
@@ -89,6 +90,16 @@ class ApplyValueSemanticsPass
       op->erase();
     }
 
+    // fix NonValueTensor type operator
+    func.walk([&](Operation* op) {
+      if (isa<OverwriteTensorContentsOp>(op)) {
+        rewriter.setInsertionPoint(op);
+        auto nonValueTensor = rewriter.create<CopyToNonValueTensorOp>(
+            op->getLoc(), op->getOperands()[1]);
+        op->setOperand(1, nonValueTensor);
+      }
+    });
+
     func.walk([&](Operation* op) {
       for (auto val : op->getResults()) {
         auto valTy = val.getType().dyn_cast_or_null<ValueTensorType>();
@@ -112,12 +123,6 @@ class ApplyValueSemanticsPass
         math::MathDialect>();
     target.addLegalOp<ModuleOp>();
 
-    target.addIllegalOp<TensorStaticInfoCastOp>();
-    target.addIllegalOp<CopyToValueTensorOp>();
-    target.addIllegalOp<CopyToNonValueTensorOp>();
-    target.addIllegalOp<ToBuiltinTensorOp>();
-    target.addIllegalOp<FromBuiltinTensorOp>();
-
     TypeConverter typeConverter;
     typeConverter.addConversion([](Type type) { return type; });
 
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
index d8b81d47f5e..c099dad9128 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
@@ -24,6 +24,7 @@
 #include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
+#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
 #include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
 using namespace mlir;
 using namespace mlir::torch;
@@ -80,7 +81,6 @@ class ConvertOperatorOp : public OpConversionPattern<OperatorOp> {
     BOOL_VAR_FROM_CONST_OPERAND(useDynamic, 9);
     BOOL_VAR_FROM_CONST_OPERAND(usePerChannel, 10);
 #undef BOOL_VAR_FROM_CONST_OPERAND
-
     auto torchMlirResultTy =
         op.getResult(0).getType().dyn_cast<ValueTensorType>();
     auto resultTy = getTypeConverter()
diff --git a/pytorch_blade/torch_blade/dynamo/__init__.py b/pytorch_blade/torch_blade/dynamo/__init__.py
index 0b263ea842f..af296cc4cd4 100644
--- a/pytorch_blade/torch_blade/dynamo/__init__.py
+++ b/pytorch_blade/torch_blade/dynamo/__init__.py
@@ -201,7 +201,8 @@ def _get_disc_decomp():
     bw_compiler=disc_compile,
     # NB: lambda here is to delay import of inductor
     decompositions=_get_disc_decomp(),
-    partition_fn=min_cut_rematerialization_partition)
+    partition_fn=min_cut_rematerialization_partition,
+    keep_inference_input_mutations=True)
 
 
 aot_disc_debug = aot_autograd(
@@ -210,7 +211,8 @@ def _get_disc_decomp():
     bw_compiler=disc_compile_ts,
     # NB: lambda here is to delay import of inductor
     decompositions=_get_disc_decomp(),
-    partition_fn=min_cut_rematerialization_partition)
+    partition_fn=min_cut_rematerialization_partition,
+    keep_inference_input_mutations=True)
 
 
 register_backend(name="disc", compiler_fn=disc)
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
index 3c25164cd20..5e4fb752a28 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -1100,6 +1100,16 @@ LogicalResult WhereOp::verify() {
   return success();
 }
 
+// ArgsMutationOp
+//===----------------------------------------------------------------------===//
+LogicalResult ArgsMutationOp::reifyReturnTypeShapes(
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
+  return success();
+}
+
+LogicalResult ArgsMutationOp::verify() { return Verify(*this); }
+
 }  // namespace mhlo_disc
 }  // namespace mlir
 
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
index 12dbe635c03..9f2bae7b65b 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
@@ -392,4 +392,18 @@ def MHLO_WhereOp : HLODISC_ShapedInterfaceOp<"where", [Pure]> {
   );
 }
 
+def MHLO_ArgsMutationOp : HLODISC_ShapedInterfaceOp<"args_mutation", []> {
+  let summary = "Argument mutation annotation operator";
+  let description = [{
+    This operator is used to annotate the mutation of arguments.
+    It is used to indicate that the arguments are mutated in the body of the
+    function. This is used to prevent the arguments from being hoisted out of
+    the function.
+  }];
+  let arguments = (ins
+    MHLO_Tensor:$input,
+    MHLO_Tensor:$output
+  );
+}
+
 #endif // MHLO_DISC_OPS
diff --git a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
index 3f8ed6e3781..aa3d100eb0e 100644
--- a/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/lhlo_disc_ops.td
@@ -282,4 +282,19 @@ def LHLODISC_ConcatenateOp : LHLODISC_Op<"concatenate", []> {
    );
 }
 
+def LHLO_ArgsMutationOp : LHLODISC_Op<"args_mutation", []> {
+  let summary = "Argument mutation annotation operator";
+  let description = [{
+    This operator is used to annotate the mutation of arguments.
+    It is used to indicate that the arguments are mutated in the body of the
+    function. This is used to prevent the arguments from being hoisted out of
+    the function.
+  }];
+  let arguments = (ins
+    Arg<LHLO_Buffer, "", [MemRead]>:$input,
+    Arg<LHLO_Buffer, "", [MemWrite]>:$output
+  );
+}
+
+
 #endif // LMHLO_DISC_OPS
diff --git a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
index 1828c032440..e29c74c2314 100644
--- a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
@@ -160,6 +160,23 @@ class HloToLhloOpConverter : public BaseOpConversion<HloOpTy> {
   }
 };
 
+struct HloToLhloArgsMutationOpConverter
+    : public BaseOpConversion<ArgsMutationOp> {
+ public:
+  using BaseOpConversion<ArgsMutationOp>::BaseOpConversion;
+  LogicalResult matchAndRewrite(
+      ArgsMutationOp hloOp, OpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    Operation* op = hloOp.getOperation();
+    auto operands = adaptor.getOperands();
+    SmallVector<Value, 4> buffer_args(operands.begin(), operands.end());
+    auto lhloOp = rewriter.create<lmhlo_disc::ArgsMutationOp>(
+        op->getLoc(), llvm::None, buffer_args);
+    rewriter.replaceOp(op, ArrayRef<Value>(buffer_args));
+    return success();
+  }
+};
+
 struct HloToLhloCustomCallOpConverter : public BaseOpConversion<CustomCallOp> {
  public:
   using BaseOpConversion<CustomCallOp>::BaseOpConversion;
@@ -256,6 +273,7 @@ struct DiscHloLegalizeToLhlo
                            tensor::TensorDialect, lmhlo::LmhloDialect>();
     target.addIllegalDialect<mhlo_disc::MhloDiscDialect>();
     target.addIllegalOp<disc_shape::TieShapeOp>();
+    target.addIllegalOp<mhlo_disc::ArgsMutationOp>();
 
     bufferization::BufferizeTypeConverter converter;
     populateDiscHLOToLHLOConversionPattern(&context, &converter, &patterns);
@@ -281,6 +299,7 @@ void populateDiscHLOToLHLOConversionPattern(
       HloToLhloOpConverter<mhlo_disc::SparseSegmentReductionOp>,
       HloToLhloOpConverter<mhlo_disc::SparseSegmentReductionWithEmptyRowsOp>,
       HloToLhloOpConverter<mhlo_disc::WhereOp>,
+      HloToLhloArgsMutationOpConverter,
       HloToLhloCustomCallOpConverter,
       HloToLhloCustomCallOpV2Converter,
       TieShapeOpConverter
diff --git a/tao_compiler/mlir/disc/transforms/disc_lhlo_rewriter.cc b/tao_compiler/mlir/disc/transforms/disc_lhlo_rewriter.cc
index d8add3a2079..605a2c4473a 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lhlo_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lhlo_rewriter.cc
@@ -57,6 +57,20 @@ namespace {
 template <typename T>
 using BaseOpConversion = OpConversionPattern<T>;
 
+struct LhloArgsMutationOpRewriter
+    : public OpRewritePattern<lmhlo_disc::ArgsMutationOp> {
+  explicit LhloArgsMutationOpRewriter(MLIRContext* context)
+      : OpRewritePattern(context) {}
+  LogicalResult matchAndRewrite(lmhlo_disc::ArgsMutationOp lhloOp,
+                                PatternRewriter& rewriter) const override {
+    auto op = lhloOp.getOperation();
+    auto operands = op->getOperands();
+    operands[0].replaceAllUsesWith(operands[1]);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 struct LhloConcatenateOpConverter
     : public OpRewritePattern<lmhlo::ConcatenateOp> {
   explicit LhloConcatenateOpConverter(MLIRContext* context)
@@ -140,6 +154,7 @@ struct DiscLhloRewriterPass
     target.addIllegalOp<lmhlo::ConcatenateOp>();
 
     patterns.insert<LhloConcatenateOpConverter>(&context);
+    patterns.insert<LhloArgsMutationOpRewriter>(&context);
     if (failed(
             applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       signalPassFailure();


Test patch:
diff --git a/pytorch_blade/tests/disc/ops/test_input_mutation.py b/pytorch_blade/tests/disc/ops/test_input_mutation.py
new file mode 100644
index 00000000000..e71fae9cc6e
--- /dev/null
+++ b/pytorch_blade/tests/disc/ops/test_input_mutation.py
@@ -0,0 +1,45 @@
+# Copyright 2023 The BladeDISC Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import torch
+from typing import List, Optional, Tuple
+from torch import Tensor
+import torch_blade
+import unittest
+from tests.disc.testing_base import skipTorchLE
+import torch_blade.clustering.support_fusion_group as fusion
+from tests.disc.testing_base import DiscTestCase
+
+class TestInputMutation(DiscTestCase):
+    def setUp(self):
+        super().setUp()
+        os.environ["TORCH_MHLO_OP_WHITE_LIST"] = "aten::copy_;aten::add"
+
+    def tearDown(self):
+        del os.environ["TORCH_MHLO_OP_WHITE_LIST"]
+
+    @skipTorchLE("2.0.0")
+    def test_input_mutation(self):
+        def func(add : Tensor, value : Tensor) -> Tensor:
+            add.add_(value)
+            return add
+        
+        with fusion.min_group_nodes(1):
+            opt_func = torch.compile(backend='aot_disc')(func)
+            add = torch.randn(8, 64, device=self.device)
+            value = torch.randn(8, 64, device=self.device)
+            actual = opt_func(add.clone(), value.clone())
+            expect = func(add.clone(), value.clone())
+            self.assertTrue(torch.allclose(actual.cpu(), expect.cpu()))
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/pytorch_blade/tests/mhlo/custom_fake_quant.mlir b/pytorch_blade/tests/mhlo/custom_fake_quant.mlir
index 521dbd66b1b..76abab94a26 100644
--- a/pytorch_blade/tests/mhlo/custom_fake_quant.mlir
+++ b/pytorch_blade/tests/mhlo/custom_fake_quant.mlir
@@ -23,4 +23,4 @@ func.func @torch_blade_custom_fake_quant(%arg0: !torch.tensor, %arg1: !torch.ten
   %5 = torch.tensor_static_info_cast %4 : !torch.tensor<[2,3],f32> to !torch.tensor
   // CHECK: return
   return %5 : !torch.tensor
-}
+}
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/tests/input-mutation.mlir b/tao_compiler/mlir/disc/transforms/tests/input-mutation.mlir
new file mode 100644
index 00000000000..b7266f1c563
--- /dev/null
+++ b/tao_compiler/mlir/disc/transforms/tests/input-mutation.mlir
@@ -0,0 +1,8 @@
+// RUN: disc-opt %s -disc-hlo-legalize-to-lhlo -hlo-legalize-to-lhlo -canonicalize -disc-lhlo-rewriter -split-input-file | FileCheck %s
+
+func.func @input_mutation(%arg0: tensor<8x32xf32>, %arg1: tensor<8x32xf32>) -> tensor<8x32xf32> {
+  // CHECK: "lmhlo.add"(%arg0, %arg1, %arg0) : (memref<8x32xf32>, memref<8x32xf32>, memref<8x32xf32>) -> ()
+  %1 = mhlo.add %arg0, %arg1 : (tensor<8x32xf32>, tensor<8x32xf32>) -> tensor<8x32xf32>
+  "mhlo_disc.args_mutation"(%1, %arg0) : (tensor<8x32xf32>, tensor<8x32xf32>) -> ()
+  "func.return"(%1) : (tensor<8x32xf32>) -> ()
+}
\ No newline at end of file