From 2f68e76bd2079cd6a18837251f4531b16b4c30fe Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 1 Jul 2026 11:30:32 +0200 Subject: [PATCH] tblgen: add ResultIndex primitive --- enzyme/Enzyme/MLIR/Implementations/Common.td | 4 ++ .../MLIR/Implementations/MathDerivatives.td | 2 +- enzyme/test/MLIR/ForwardMode/exp_wrap.mlir | 21 +++++++ enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp | 59 +++++++++++++++++++ 4 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 enzyme/test/MLIR/ForwardMode/exp_wrap.mlir diff --git a/enzyme/Enzyme/MLIR/Implementations/Common.td b/enzyme/Enzyme/MLIR/Implementations/Common.td index 01a539695100..827e628c904e 100644 --- a/enzyme/Enzyme/MLIR/Implementations/Common.td +++ b/enzyme/Enzyme/MLIR/Implementations/Common.td @@ -81,6 +81,10 @@ class Operation { bit usesCustom = usesCustom_; } +class ResultIndex { + int index = index_; +} +def Result : ResultIndex; class DiffeRetIndex indices_> { list indices = indices_; } diff --git a/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td b/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td index 5905542eb212..aa468c52540c 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td +++ b/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td @@ -7,7 +7,7 @@ def : MLIRDerivative<"math", "CosOp", (Op $x), >; def : MLIRDerivative<"math", "ExpOp", (Op $x), [ - (CheckedMulF (DiffeRet), (ExpF $x)) + (CheckedMulF (DiffeRet), (Result)) ] >; def : MLIRDerivative<"math", "SinOp", (Op $x), diff --git a/enzyme/test/MLIR/ForwardMode/exp_wrap.mlir b/enzyme/test/MLIR/ForwardMode/exp_wrap.mlir new file mode 100644 index 000000000000..1112c6f3d51c --- /dev/null +++ b/enzyme/test/MLIR/ForwardMode/exp_wrap.mlir @@ -0,0 +1,21 @@ +// RUN: %eopt --enzyme-wrap="infn=main outfn= argTys=enzyme_dup retTys=enzyme_dup mode=ForwardMode" --canonicalize --remove-unnecessary-enzyme-ops %s | FileCheck %s --check-prefix=FWD +// RUN: %eopt --enzyme-wrap="infn=main outfn= argTys=enzyme_active retTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --enzyme-simplify-math %s | FileCheck %s --check-prefix=REV + +module { + func.func @main(%arg0: f32) -> f32 { + %0 = math.exp %arg0 : f32 + return %0 : f32 + } +} + +// FWD: func.func @main(%[[x:.+]]: f32, %[[dx:.+]]: f32) -> (f32, f32) { +// FWD-NEXT: %[[exp:.+]] = math.exp %[[x]] : f32 +// FWD-NEXT: %[[shadow:.+]] = arith.mulf %[[dx]], %[[exp]] : f32 +// FWD-NEXT: return %[[exp]], %[[shadow]] : f32, f32 +// FWD-NEXT: } + +// REV: func.func @main(%[[x:.+]]: f32, %[[dret:.+]]: f32) -> f32 { +// REV-NEXT: %[[exp:.+]] = math.exp %[[x]] : f32 +// REV-NEXT: %[[darg:.+]] = arith.mulf %[[dret]], %[[exp]] : f32 +// REV-NEXT: return %[[darg]] : f32 +// REV-NEXT: } diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp index b7e3ed274072..c00062ce4b10 100644 --- a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp @@ -432,6 +432,47 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, os << curIndent << INDENT << "out; })\n"; } return true; + } else if (opName == "ResultIndex" || Def->isSubClassOf("ResultIndex")) { + if (intrinsic != MLIRDerivatives) { + PrintFatalError(pattern->getLoc(), + "ResultIndex only supported in MLIR derivatives"); + return false; + } + + if (resultRoot->getNumArgs() != 0) { + PrintFatalError(pattern->getLoc(), "no arg supported for ResultIndex"); + return false; + } + + auto index = dyn_cast(Def->getValueInit("index")); + + os << "({" + << " // Computing ResultIndex (index=" << index->getValue() << ")" + << argPattern << "\n"; + + llvm::SmallString<32> buf; + if (startsWith(argPattern.toStringRef(buf), "fwdarg")) { + os << curIndent << INDENT + << "mlir::Value resultIndex = " + "gutils->getNewFromOriginal(op->getResult(" + << index->getValue() << "));\n"; + os << curIndent << INDENT + << "builder.setInsertionPointAfterValue(resultIndex)" + ";\n"; + os << curIndent << INDENT << "resultIndex;\n"; + } else { + os << curIndent << INDENT + << "auto neededArgs = cachedArguments(op, gutils);\n"; + os << curIndent << INDENT << "int numNeededArgs = 0;\n"; + os << curIndent << INDENT << "for (auto needed : neededArgs)\n"; + os << curIndent << INDENT << INDENT + << "if (needed) numNeededArgs += 1;\n"; + os << curIndent << INDENT << "gutils->popCache(caches[numNeededArgs + " + << index->getValue() << "], builder);\n"; + } + os << curIndent << "})"; + + return false; } else if (opName == "TypeOf" || Def->isSubClassOf("TypeOf")) { if (resultRoot->getNumArgs() != 1) PrintFatalError(pattern->getLoc(), "only single op TypeOf supported"); @@ -1408,6 +1449,8 @@ void handleUse( if (opName == "InactiveArgSpec" || Def->isSubClassOf("InactiveArgSpec")) { return; } + if (opName == "ResultIndex" || Def->isSubClassOf("ResultIndex")) + return; if (!Def->isSubClassOf("Operation")) { errs() << *resultTree << "\n"; errs() << opName << " " << *Def << "\n"; @@ -1728,6 +1771,13 @@ static void emitMLIRReverse(raw_ostream &os, const Record *pattern, os << " return toret;\n"; os << " }\n"; + os << " SmallVector cachedResults(Operation *op,\n" + " MGradientUtilsReverse *gutils) " + "const {\n" + " SmallVector neededResults(op->getNumResults(), " + "false);\n" + " return neededResults;\n}\n"; + os << " SmallVector cacheValues(Operation *op,\n"; os << " MGradientUtilsReverse *gutils) " "const {\n"; @@ -1743,6 +1793,15 @@ static void emitMLIRReverse(raw_ostream &os, const Record *pattern, "getOperand(en.index())), builder);\n"; os << " toret.push_back(cache);\n"; os << " }\n"; + os << " " + "builder.setInsertionPointAfter(gutils->getNewFromOriginal(op));\n"; + os << " auto neededResults = cachedResults(op, gutils);\n"; + os << " for (auto en : llvm::enumerate(neededResults)) {\n"; + os << " Value cache = " + "gutils->initAndPushCache(gutils->getNewFromOriginal(op->getResult(en." + "index())), builder);\n"; + os << " toret.push_back(cache);\n"; + os << " }"; os << " return toret;\n"; os << " }\n"; os << "\n";