Add element preserving trait to optimize involution folding

[ahmedsabie]

Summary

Adding this trait allows involution folding to be applicable in more situations

[rriddle]

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

[ahmedsabie]

In D89572#2335343, @rriddle wrote:

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

This makes the original trait applicable in more places because it can ignore nodes that don't actually end up changing tensor values so you can still fold. This is currently used in tensorflow grappler optimizations under the same name

[rriddle]

In D89572#2337634, @ahmedsabie wrote:

In D89572#2335343, @rriddle wrote:

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

This makes the original trait applicable in more places because it can ignore nodes that don't actually end up changing tensor values so you can still fold.

I don't understand this aspect. Is this supposed to mean that the operation would somehow mutate the input tensors? Putting the fact that tensors are value-typed aside, I would expect this to be something covered by side effects and not a trait like this. If an operation is going to mutate its input operands, that is something that should be expressed as a side effect.

This is currently used in tensorflow grappler optimizations under the same name

I would prefer to always use names that make sense in the context of MLIR, and not carry over names just because they are used in a different infrastructure. If a name can be carried over that's great(e.g. if it is widely used throughout the domain), but in this case it just doesn't make sense given how general of a name it is.

[ahmedsabie]

In D89572#2337635, @rriddle wrote:

In D89572#2337634, @ahmedsabie wrote:

In D89572#2335343, @rriddle wrote:

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

This makes the original trait applicable in more places because it can ignore nodes that don't actually end up changing tensor values so you can still fold.

I don't understand this aspect. Is this supposed to mean that the operation would somehow mutate the input tensors? Putting the fact that tensors are value-typed aside, I would expect this to be something covered by side effects and not a trait like this. If an operation is going to mutate its input operands, that is something that should be expressed as a side effect.
No I mean if you had abs(permute(abs(x))) then that's the same as permute(abs(x)) but you can't generally apply this optimization to all side effect free operations, example abs(negate(abs(x))) is not the same as negate(abs(x)). ValuePreserving is capturing the property of permute that it doesn't actually produce new values when applied

[rriddle]

In D89572#2337652, @ahmedsabie wrote:

In D89572#2337635, @rriddle wrote:

In D89572#2337634, @ahmedsabie wrote:

In D89572#2335343, @rriddle wrote:

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

This makes the original trait applicable in more places because it can ignore nodes that don't actually end up changing tensor values so you can still fold.

I don't understand this aspect. Is this supposed to mean that the operation would somehow mutate the input tensors? Putting the fact that tensors are value-typed aside, I would expect this to be something covered by side effects and not a trait like this. If an operation is going to mutate its input operands, that is something that should be expressed as a side effect.

No I mean if you had abs(permute(abs(x))) then that's the same as permute(abs(x)) but you can't generally apply this optimization to all side effect free operations, example abs(negate(abs(x))) is not the same as negate(abs(x)). ValuePreserving is capturing the property of permute that it doesn't actually produce new values when applied

Thank you for explaining, I understand what you mean by value preserving now. That being said, I don't think this optimization is legal with just idempotent and "value preserving". For example, stable_sort(permute(stable_sort(x))) is not strictly the same as permute(stable_sort(x)).

[ahmedsabie]

Update name of trait and apply it to involutions instead

[andyly]

mlir/include/mlir/IR/OpDefinition.h

1061–1062

Can this comment be updated to not be the same as IsIdempotent?

[ahmedsabie]

Update comment

[ahmedsabie]

In D89572#2337791, @rriddle wrote:

In D89572#2337652, @ahmedsabie wrote:

In D89572#2337635, @rriddle wrote:

In D89572#2337634, @ahmedsabie wrote:

In D89572#2335343, @rriddle wrote:

I don't really understand the motivation here, can you detail it a bit more. Why is the original trait not enough? Also, ValuePreserving seems too broad of a name for what this is doing.

This makes the original trait applicable in more places because it can ignore nodes that don't actually end up changing tensor values so you can still fold.

I don't understand this aspect. Is this supposed to mean that the operation would somehow mutate the input tensors? Putting the fact that tensors are value-typed aside, I would expect this to be something covered by side effects and not a trait like this. If an operation is going to mutate its input operands, that is something that should be expressed as a side effect.

No I mean if you had abs(permute(abs(x))) then that's the same as permute(abs(x)) but you can't generally apply this optimization to all side effect free operations, example abs(negate(abs(x))) is not the same as negate(abs(x)). ValuePreserving is capturing the property of permute that it doesn't actually produce new values when applied

Thank you for explaining, I understand what you mean by value preserving now. That being said, I don't think this optimization is legal with just idempotent and "value preserving". For example, stable_sort(permute(stable_sort(x))) is not strictly the same as permute(stable_sort(x)).

That's a good point, it would only work if we define idempotent to split across elements but that wouldn't be called idempotent necessarily anymore. I updated the diff with a similar idea but for involution which should work although I added it as a pass rather than a trait fold because it can't be implemented as a fold I believe

[rriddle]

Can you formulate a formal description for what you consider "element preserving"? I still have trouble understanding the implications just by looking at the name and the equation notation.

mlir/include/mlir/IR/OpBase.td

1727

nit: Use == here?

mlir/include/mlir/Transforms/Trait.h

4

I would not put this in the Transforms/ library because it prevents dialects from depending on it, i.e dialect libraries should never really depend on Transforms/.

mlir/lib/Transforms/Trait.cpp

9

Can you add a detailed comment of the folding here?

[rriddle]

mlir/include/mlir/Transforms/Trait.h

4

We already have some common foldings near the trait definitions, this could go there for now.

[ahmedsabie]

mlir/include/mlir/Transforms/Trait.h

4

That would be the OpDefinition.h file right? The issue is this requires inheriting from TraitRewritePattern which is in PatternMatch.h and PatternMatch.h itself relies on OpDefinition.h

[rriddle]

mlir/include/mlir/IR/PatternMatch.h

234

I think checking for a trait is easy enough to do inside of a normal rewrite pattern that this is probably overkill.

mlir/include/mlir/Transforms/Trait.h

4

Yes, Normally we expose patterns via some form of "populate*Patterns" method. Using that technique would remove the need for RewritePattern to be defined in the .h file. populateInvolutionCanonicalizationPatterns?

Orthogonally, this case starts to make me think that we should maybe have a 'getCanonicalizationPatterns' hook for traits as well, so this gets auto exposed?

mlir/include/mlir/IR/OpBase.td

 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // op op X == op X
 def Idempotent  : NativeOpTrait<"IsIdempotent">;
 // op op X == X
 def Involution  : NativeOpTrait<"IsInvolution">;
+// op op' X = op' op X for all per-element unary op'.
+def ElementPreserving  : NativeOpTrait<"IsElementPreserving">;
 // Op behaves like a constant.
 def ConstantLike : NativeOpTrait<"ConstantLike">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;

mlir/include/mlir/IR/OpDefinition.h

 namespace impl {
 OpFoldResult foldIdempotent(Operation *op);
 OpFoldResult foldInvolution(Operation *op);
 LogicalResult verifyZeroOperands(Operation *op);
 LogicalResult verifyOneOperand(Operation *op);
 LogicalResult verifyNOperands(Operation *op, unsigned numOperands);
 LogicalResult verifyIsIdempotent(Operation *op);
 LogicalResult verifyIsInvolution(Operation *op);
+LogicalResult verifyIsElementPreserving(Operation *op);
 LogicalResult verifyAtLeastNOperands(Operation *op, unsigned numOperands);
 LogicalResult verifyOperandsAreFloatLike(Operation *op);
 LogicalResult verifyOperandsAreSignlessIntegerLike(Operation *op);
 LogicalResult verifySameTypeOperands(Operation *op);
 LogicalResult verifyZeroRegion(Operation *op);
 LogicalResult verifyOneRegion(Operation *op);
 LogicalResult verifyNRegions(Operation *op, unsigned numRegions);
 LogicalResult verifyAtLeastNRegions(Operation *op, unsigned numRegions);
static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyIsIdempotent(op);
   }
 
   static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
     return impl::foldIdempotent(op);
   }
 };
 
+/// This class adds property that the operation is element preserving.
+/// This means a unary to unary operation "f" that satisfies
+/// f(g(x)) = g(f(x)) for all per element unary ops g
+template <typename ConcreteType>
+class IsElementPreserving
+    : public TraitBase<ConcreteType, IsElementPreserving> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<OneResult>(),
+                  "expected operation to produce one result");
+    static_assert(ConcreteType::template hasTrait<OneOperand>(),
+                  "expected operation to take one operand");
+    static_assert(ConcreteType::template hasTrait<SameOperandsAndResultType>(),
+                  "expected operation to preserve type");
+    // ElementPreserving requires the operation to be side effect free as well
+    // but currently this check is under a FIXME and is not actually done.
+    return impl::verifyIsElementPreserving(op);
+  }
+};
+
 /// This class verifies that all operands of the specified op have a float type,
 /// a vector thereof, or a tensor thereof.
 template <typename ConcreteType>
 class OperandsAreFloatLike
     : public TraitBase<ConcreteType, OperandsAreFloatLike> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifyOperandsAreFloatLike(op);

mlir/include/mlir/IR/PatternMatch.h

virtual LogicalResult matchAndRewrite(SourceOp op,
     if (succeeded(match(op))) {
       rewrite(op, rewriter);
       return success();
     }
     return failure();
   }
 };
 
+/// TraitRewritePattern is a wrapper around RewritePattern that allows for
+/// matching and rewriting against an op that matches a trait
+template <template <typename> class Trait, typename TraitPattern>
+struct TraitRewritePattern : public RewritePattern {
+  /// Patterns can specify the benefit of the pattern matching.
+  TraitRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
+      : RewritePattern({}, benefit, context, MatchAnyOpTypeTag{}) {}
+
+  /// Wrappers around the RewritePattern methods that pass the op if it
+  /// matches the trait.
+  void rewrite(Operation *op, PatternRewriter &rewriter) const final {
+    static_cast<const TraitPattern *>(this)->rewriteTrait(op, rewriter);
+  }
+  LogicalResult match(Operation *op) const final {
+    if (op->hasTrait<Trait>())
+      return static_cast<const TraitPattern *>(this)->matchTrait(op);
+    return failure();
+  }
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const final {
+    if (op->hasTrait<Trait>())
+      return static_cast<const TraitPattern *>(this)->matchAndRewriteTrait(
+          op, rewriter);
+    return failure();
+  }
+
+  /// Rewrite and Match methods that operate on the op with the given
+  /// trait. These must be overridden by the derived pattern class.
+  void rewriteTrait(Operation *op, PatternRewriter &rewriter) const {
+    llvm_unreachable("must override rewrite or matchAndRewrite");
+  }
+  LogicalResult matchTrait(Operation *op) const {
+    llvm_unreachable("must override match or matchAndRewrite");
+  }
+  LogicalResult matchAndRewriteTrait(Operation *op,
+                                     PatternRewriter &rewriter) const {
+    if (succeeded(matchTrait(op))) {
+      rewriteTrait(op, rewriter);
+      return success();
+    }
+    return failure();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // PatternRewriter class
 //===----------------------------------------------------------------------===//
 
 /// This class coordinates the application of a pattern to the current function,
 /// providing a way to create operations and keep track of what gets deleted.
 ///
 /// These class serves two purposes:

mlir/include/mlir/Transforms/Trait.h

+#ifndef MLIR_TRANSFORMS_TRAIT_H
+#define MLIR_TRANSFORMS_TRAIT_H
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+/// Folds Involution trait labelled ops across ElementPreserving labelled ops.
+class FoldInvolutionsThroughElementPreserving
+    : public TraitRewritePattern<OpTrait::IsInvolution,
+                                 FoldInvolutionsThroughElementPreserving> {
+public:
+  using TraitRewritePattern<
+      OpTrait::IsInvolution,
+      FoldInvolutionsThroughElementPreserving>::TraitRewritePattern;
+  /// Performs the actual rewriting step.
+  LogicalResult matchAndRewriteTrait(Operation *op,
+                                     PatternRewriter &rewriter) const;
+};
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_TRAIT_H
 No newline at end of file

mlir/lib/IR/Operation.cpp

 LogicalResult OpTrait::impl::verifyIsInvolution(Operation *op) {
   // FIXME: Add back check for no side effects on operation.
   // Currently adding it would cause the shared library build
   // to fail since there would be a dependency of IR on SideEffectInterfaces
   // which is cyclical.
   return success();
 }
 
+LogicalResult OpTrait::impl::verifyIsElementPreserving(Operation *op) {
+  // FIXME: Add back check for no side effects on operation.
+  // Currently adding it would cause the shared library build
+  // to fail since there would be a dependency of IR on SideEffectInterfaces
+  // which is cyclical.
+  return success();
+}
+
 LogicalResult
 OpTrait::impl::verifyOperandsAreSignlessIntegerLike(Operation *op) {
   for (auto opType : op->getOperandTypes()) {
     auto type = getTensorOrVectorElementType(opType);
     if (!type.isSignlessIntOrIndex())
       return op->emitOpError() << "requires an integer or index type";
   }
   return success();

mlir/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTransforms
   MemRefDataFlowOpt.cpp
   NormalizeMemRefs.cpp
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
+  Trait.cpp
   ViewOpGraph.cpp
   ViewRegionGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRStandardOpsIncGen

mlir/lib/Transforms/Trait.cpp

+#include "mlir/Transforms/Trait.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// FoldInvolutionsThroughElement
+//===----------------------------------------------------------------------===//
+
+LogicalResult FoldInvolutionsThroughElementPreserving::matchAndRewriteTrait(
+    Operation *op, PatternRewriter &rewriter) const {
+
+  auto *argumentOp = op->getOperand(0).getDefiningOp();
+
+  if (argumentOp && op->getName() == argumentOp->getName())
+    return failure();
+
+  while (argumentOp && op->getName() != argumentOp->getName() &&
+         argumentOp->hasTrait<OpTrait::IsElementPreserving>() &&
+         argumentOp->hasOneUse()) {
+    // Involution and value preserving ops all have one arg.
+    assert(argumentOp->getNumOperands() == 1);
+    argumentOp = argumentOp->getOperand(0).getDefiningOp();
+  }
+
+  if (argumentOp && op->getName() == argumentOp->getName() &&
+      argumentOp->hasOneUse()) {
+    rewriter.replaceOp(argumentOp, {argumentOp->getOperand(0)});
+    rewriter.replaceOp(op, {op->getOperand(0)});
+    return success();
+  }
+
+  return failure();
+}
 No newline at end of file

mlir/test/lib/Dialect/Test/TestOps.td

 def TestInvolutionTraitSuccesfulOperationFolderOp
  : TEST_Op<"op_involution_trait_succesful_operation_fold",
            [SameOperandsAndResultType, NoSideEffect, Involution]> {
   let arguments = (ins I32:$op1);
   let results = (outs I32);
   let hasFolder = 1;
 }
 
+def TestElementPreservingTraitOp
+ : TEST_Op<"op_element_preserving_trait",
+           [SameOperandsAndResultType, NoSideEffect, ElementPreserving]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+}
+
+def TestNoElementPreservingTraitOp
+ : TEST_Op<"op_no_element_preserving_trait",
+           [SameOperandsAndResultType, NoSideEffect]> {
+  let arguments = (ins I32:$op1);
+  let results = (outs I32);
+}
+
 def TestOpInPlaceFoldAnchor : TEST_Op<"op_in_place_fold_anchor"> {
   let arguments = (ins I32);
   let results = (outs I32);
 }
 
 def TestOpInPlaceFold : TEST_Op<"op_in_place_fold"> {
   let arguments = (ins I32:$op, I32Attr:$attr);
   let results = (outs I32);

mlir/test/lib/Dialect/Test/TestTraits.cpp

 //===- TestTraits.cpp - Test trait folding --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
+#include "mlir/Transforms/Trait.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Trait Folder.
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TestInvolutionTraitFailingOperationFolderOp::fold(
     ArrayRef<Attribute> operands) {
   // This failure should cause the trait fold to run instead.
   return {};
 }
 
 OpFoldResult TestInvolutionTraitSuccesfulOperationFolderOp::fold(
     ArrayRef<Attribute> operands) {
   auto argument_op = getOperand();
   // The success case should cause the trait fold to be supressed.
   return argument_op.getDefiningOp() ? argument_op : OpFoldResult{};
 }
 
 namespace {
 struct TestTraitFolder : public PassWrapper<TestTraitFolder, FunctionPass> {
   void runOnFunction() override {
-    applyPatternsAndFoldGreedily(getFunction(), {});
+    OwningRewritePatternList patterns;
+    patterns.insert<FoldInvolutionsThroughElementPreserving>(&getContext());
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestTraitsPass() {
   PassRegistration<TestTraitFolder>("test-trait-folder", "Run trait folding");
 }
 } // namespace mlir

mlir/test/mlir-tblgen/trait.mlir

 func @testTripleIdempotent(%arg0: i32) -> i32 {
   // CHECK: [[IDEMPOTENT:%.+]] = "test.op_idempotent_trait"([[ARG0]])
   %0 = "test.op_idempotent_trait"(%arg0) : (i32) -> i32
   %1 = "test.op_idempotent_trait"(%0) : (i32) -> i32
   %2 = "test.op_idempotent_trait"(%1) : (i32) -> i32
   // CHECK: return [[IDEMPOTENT]]
   return %2: i32
 }
+
+//===----------------------------------------------------------------------===//
+// Test that folding involution ops mixed with element preserving ops
+// works correctly
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @testSingleElementPreserving
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testSingleElementPreserving(%arg0 : i32) -> i32 {
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES:%.+]] = "test.op_element_preserving_trait"([[ARG0]])
+  %1 = "test.op_element_preserving_trait"(%0) : (i32) -> i32
+  %2 = "test.op_involution_trait_no_operation_fold"(%1) : (i32) -> i32
+  // CHECK: return [[ELEMENTPRES]]
+  return %2: i32
+}
+
+// CHECK-LABEL: func @testChainElementPreserving
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testChainElementPreserving(%arg0 : i32) -> i32 {
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES1:%.+]] = "test.op_element_preserving_trait"([[ARG0]])
+  %1 = "test.op_element_preserving_trait"(%0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES2:%.+]] = "test.op_element_preserving_trait"([[ELEMENTPRES1]])
+  %2 = "test.op_element_preserving_trait"(%1) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES3:%.+]] = "test.op_element_preserving_trait"([[ELEMENTPRES2]])
+  %3 = "test.op_element_preserving_trait"(%2) : (i32) -> i32
+  %4 = "test.op_involution_trait_no_operation_fold"(%3) : (i32) -> i32
+  // CHECK: return [[ELEMENTPRES3]]
+  return %4: i32
+}
+
+// CHECK-LABEL: func @testNoElementPreserving
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testNoElementPreserving(%arg0 : i32) -> i32 {
+  // CHECK: [[INVOLUTION1:%.+]] = "test.op_involution_trait_no_operation_fold"([[ARG0]])
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES1:%.+]] = "test.op_element_preserving_trait"([[INVOLUTION1]])
+  %1 = "test.op_element_preserving_trait"(%0) : (i32) -> i32
+  // CHECK: [[NOELEMENTPRES:%.+]] = "test.op_no_element_preserving_trait"([[ELEMENTPRES1]])
+  %2 = "test.op_no_element_preserving_trait"(%1) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES2:%.+]] = "test.op_element_preserving_trait"([[NOELEMENTPRES]])
+  %3 = "test.op_element_preserving_trait"(%2) : (i32) -> i32
+  // CHECK: return [[ELEMENTPRES2]]
+  return %3: i32
+}
+
+// CHECK-LABEL: func @testNoOneUse
+// CHECK-SAME:  ([[ARG0:%.+]]: i32)
+func @testNoOneUse(%arg0 : i32) -> (i32, i32) {
+  // CHECK: [[INVOLUTION1:%.+]] = "test.op_involution_trait_no_operation_fold"([[ARG0]])
+  %0 = "test.op_involution_trait_no_operation_fold"(%arg0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES1:%.+]] = "test.op_element_preserving_trait"([[INVOLUTION1]])
+  %1 = "test.op_element_preserving_trait"(%0) : (i32) -> i32
+  // CHECK: [[ELEMENTPRES2:%.+]] = "test.op_element_preserving_trait"([[INVOLUTION1]])
+  %2 = "test.op_element_preserving_trait"(%0) : (i32) -> i32
+  // CHECK: [[INVOLUTION2:%.+]] = "test.op_involution_trait_no_operation_fold"([[ELEMENTPRES1]])
+  %3 = "test.op_involution_trait_no_operation_fold"(%1) : (i32) -> i32
+  // CHECK: return [[ELEMENTPRES2]], [[INVOLUTION2]]
+  return %2, %3 : i32, i32
+}