diff --git a/src/enzyme_ad/jax/Dialect/Tessera/Ops.cpp b/src/enzyme_ad/jax/Dialect/Tessera/Ops.cpp index c2da6848f5..4535aa9f3b 100644 --- a/src/enzyme_ad/jax/Dialect/Tessera/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Tessera/Ops.cpp @@ -231,6 +231,89 @@ void CallOp::getEffects( effects.emplace_back(MemoryEffects::Effect::get()); } +// 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 { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(CallOp op, + PatternRewriter &rewriter) const override { + SmallVector newOperands; + bool changed = false; + + for (Value arg : op.getOperands()) { + if (!isa(arg.getType())) { + newOperands.push_back(arg); + continue; + } + + // trace through bitcast + auto bitcast = arg.getDefiningOp(); + if (!bitcast) { + newOperands.push_back(arg); + continue; + } + + // trace through insertelement chain, following vector operand until we + // hit poison + Value vec = bitcast.getArg(); + SmallVector elements; + while (auto insertOp = vec.getDefiningOp()) { + elements.push_back(insertOp.getValue()); + vec = insertOp.getVector(); + } + if (!vec.getDefiningOp()) { + 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(); + 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(); + if (sourceCall) { + newOperands.push_back(sourceCall.getResult(0)); + changed = true; + } else { + newOperands.push_back(arg); + } + } + + if (!changed) + return failure(); + rewriter.replaceOpWithNewOp(op, op.getResultTypes(), newOperands, + op->getAttrs()); + return success(); + } +}; + +void CallOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.insert(context); +} + //===----------------------------------------------------------------------===// // ReturnOp //===----------------------------------------------------------------------===// diff --git a/src/enzyme_ad/jax/Dialect/Tessera/Ops.td b/src/enzyme_ad/jax/Dialect/Tessera/Ops.td index eb01103c19..283ff3a76a 100644 --- a/src/enzyme_ad/jax/Dialect/Tessera/Ops.td +++ b/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(); diff --git a/test/lit_tests/tessera/fold_tessera_call_chain.mlir b/test/lit_tests/tessera/fold_tessera_call_chain.mlir new file mode 100644 index 0000000000..b3e9283b06 --- /dev/null +++ b/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, byRefArgs = array, 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: } \ No newline at end of file