Canonicalize tessera calls

src/enzyme_ad/jax/Dialect/Tessera/Ops.cpp

@@ -231,6 +231,89 @@ void CallOp::getEffects(
   effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Free>());
 }
 
+// Folds the extractvalue/insertelement/bitcast chain that appears between two
+// tessera.call ops when the struct result of the first call is repacked into an
+// integer to be passed as an argument to the second call. Replaces the chain
+// with a direct use of the struct value.
+class FoldTesseraCallChain final : public OpRewritePattern<CallOp> {
+public:
+  using OpRewritePattern<CallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CallOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> newOperands;
+    bool changed = false;
+
+    for (Value arg : op.getOperands()) {
+      if (!isa<IntegerType>(arg.getType())) {
+        newOperands.push_back(arg);
+        continue;
+      }
+
+      // trace through bitcast
+      auto bitcast = arg.getDefiningOp<LLVM::BitcastOp>();
+      if (!bitcast) {
+        newOperands.push_back(arg);
+        continue;
+      }
+
+      // trace through insertelement chain, following vector operand until we
+      // hit poison
+      Value vec = bitcast.getArg();
+      SmallVector<Value> elements;
+      while (auto insertOp = vec.getDefiningOp<LLVM::InsertElementOp>()) {
+        elements.push_back(insertOp.getValue());
+        vec = insertOp.getVector();
+      }
+      if (!vec.getDefiningOp<LLVM::PoisonOp>()) {
+        newOperands.push_back(arg);
+        continue;
+      }
+
+      // trace through extractvalue chain and make sure all come from the same
+      // source call
+      Value source;
+      for (Value elem : elements) {
+        auto extractOp = elem.getDefiningOp<LLVM::ExtractValueOp>();
+        if (!extractOp) {
+          source = nullptr;
+          break;
+        }
+        if (!source)
+          source = extractOp.getContainer();
+        else if (source != extractOp.getContainer()) {
+          source = nullptr; // elements came from different sources
+          break;
+        }
+      }
+      if (!source) {
+        newOperands.push_back(arg);
+        continue;
+      }
+
+      // check source is a tessera call
+      auto sourceCall = source.getDefiningOp<CallOp>();
+      if (sourceCall) {
+        newOperands.push_back(sourceCall.getResult(0));
+        changed = true;
+      } else {
+        newOperands.push_back(arg);
+      }
+    }
+
+    if (!changed)
+      return failure();
+    rewriter.replaceOpWithNewOp<CallOp>(op, op.getResultTypes(), newOperands,
+                                        op->getAttrs());
+    return success();
+  }
+};
+
+void CallOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.insert<FoldTesseraCallChain>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//

src/enzyme_ad/jax/Dialect/Tessera/Ops.td

@@ -137,6 +137,8 @@ def CallOp : TesseraOp<"call",
             results, operands);
     }]>];
 
+  let hasCanonicalizer = 1;
+
   let extraClassDeclaration = [{
     FunctionType getCalleeType();
 

test/lit_tests/tessera/fold_tessera_call_chain.mlir

@@ -0,0 +1,36 @@
+// RUN: enzymexlamlir-opt %s -canonicalize | FileCheck %s
+
+module {
+  tessera.define @eigen.inv(%arg0: !llvm.ptr {llvm.sret = !llvm.struct<(array<4 x f32>)>, llvm.nocapture, llvm.writeonly}, %arg1: !llvm.ptr {llvm.nocapture, llvm.readonly}) -> () attributes {argSizes = array<i64: 16, 16>, byRefArgs = array<i1: true, true>, pure = true} {
+    tessera.return
+  }
+
+  llvm.func @main(%arg0: i128) -> !llvm.struct<(array<4 x f32>)> {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    %1 = llvm.mlir.constant(1 : i32) : i32
+    %2 = llvm.mlir.constant(2 : i32) : i32
+    %3 = llvm.mlir.constant(3 : i32) : i32
+    %4 = llvm.mlir.poison : vector<4xf32>
+    %5 = tessera.call @eigen.inv(%arg0) : (i128) -> !llvm.struct<(array<4 x f32>)>
+    %6 = llvm.extractvalue %5[0, 0] : !llvm.struct<(array<4 x f32>)> 
+    %7 = llvm.extractvalue %5[0, 1] : !llvm.struct<(array<4 x f32>)> 
+    %8 = llvm.extractvalue %5[0, 2] : !llvm.struct<(array<4 x f32>)>
+    %9 = llvm.extractvalue %5[0, 3] : !llvm.struct<(array<4 x f32>)>
+    %10 = llvm.insertelement %6, %4[%0 : i32] : vector<4xf32>
+    %11 = llvm.insertelement %7, %10[%1 : i32] : vector<4xf32> 
+    %12 = llvm.insertelement %8, %11[%2 : i32] : vector<4xf32>
+    %13 = llvm.insertelement %9, %12[%3 : i32] : vector<4xf32>
+    %14 = llvm.bitcast %13 : vector<4xf32> to i128
+    %15 = tessera.call @eigen.inv(%14) : (i128) -> !llvm.struct<(array<4 x f32>)>
+    llvm.return %15 : !llvm.struct<(array<4 x f32>)>
+  }
+}
+
+// CHECK: tessera.define @eigen.inv(%[[ARG0:.*]]: !llvm.ptr {llvm.nocapture, llvm.sret = !llvm.struct<(array<4 x f32>)>, llvm.writeonly}, %[[ARG1:.*]]: !llvm.ptr {llvm.nocapture, llvm.readonly})
+// CHECK-NEXT: tessera.return
+
+// CHECK: llvm.func @main(%[[VAL:.*]]: i128) -> !llvm.struct<(array<4 x f32>)> {
+// CHECK-NEXT: %[[RES1:.*]] = tessera.call @eigen.inv(%[[VAL]]) : (i128) -> !llvm.struct<(array<4 x f32>)>
+// CHECK-NEXT: %[[RES2:.*]] = tessera.call @eigen.inv(%[[RES1]]) : (!llvm.struct<(array<4 x f32>)>) -> !llvm.struct<(array<4 x f32>)>
+// CHECK-NEXT: llvm.return %[[RES2]] : !llvm.struct<(array<4 x f32>)>
+// CHECK-NEXT: }