From 562bb306f99a261d1a27ff8a2bb33e752fe6f0e4 Mon Sep 17 00:00:00 2001 From: jacobleft <88961633+jacobleft@users.noreply.github.com> Date: Sat, 18 Jul 2026 10:16:41 +0800 Subject: [PATCH 1/3] Fix vectorized MLIR derivative constants --- enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp index b7e3ed27407..09e69dce6f7 100644 --- a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp @@ -618,7 +618,12 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, std::string ord; bool shadowType = false; if (resultRoot->getNumArgs() == 0) { - ord = "op->getResult(0)"; + ord = "gutils->getShadowType(op->getResult(0)"; + // This constant participates in a tangent expression. In vector + // forward mode, the tangent has a leading width dimension even when + // the primal result does not. Use the shadow type so a constant + // such as `1 - tanh(x)^2` composes with width-vectorized operands. + shadowType = true; } else { if (resultRoot->getArgName(0)) { auto name = resultRoot->getArgName(0)->getAsUnquotedString(); From a98cc16d470ea258ee278fa0002ab2e7bfb79ac3 Mon Sep 17 00:00:00 2001 From: jacobleft <88961633+jacobleft@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:09:07 +0800 Subject: [PATCH 2/3] Fix vectorized MLIR derivative constants --- enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp index 09e69dce6f7..66271c045ff 100644 --- a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp +++ b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp @@ -608,6 +608,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, PrintFatalError(pattern->getLoc(), Twine("'value' not defined in ") + resultTree->getAsString()); + bool vectorizedConstant = false; if (intrinsic == MLIRDerivatives) { if (resultRoot->getNumArgs() > 1) PrintFatalError(pattern->getLoc(), @@ -621,8 +622,8 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, ord = "gutils->getShadowType(op->getResult(0)"; // This constant participates in a tangent expression. In vector // forward mode, the tangent has a leading width dimension even when - // the primal result does not. Use the shadow type so a constant - // such as `1 - tanh(x)^2` composes with width-vectorized operands. + // the primal result does not. Use the shadow type to preserve the + // tangent shape. shadowType = true; } else { if (resultRoot->getArgName(0)) { @@ -670,6 +671,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, os << ")"; os << ", "; os << "\"" << value->getValue() << "\"))"; + vectorizedConstant = shadowType; } else { if (resultRoot->getNumArgs() != 1) PrintFatalError(pattern->getLoc(), @@ -700,7 +702,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os, resultTree->getAsString()); os << ", \"" << value->getValue() << "\")"; } - return false; + return vectorizedConstant; } else if (opName == "Zero" || Def->isSubClassOf("Zero")) { if (resultRoot->getNumArgs() != 1) PrintFatalError(pattern->getLoc(), "only single op Zero supported"); From 448136afce908952aed0c3fe7ad5009f5f0fc52f Mon Sep 17 00:00:00 2001 From: jacobleft <88961633+jacobleft@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:59:25 +0800 Subject: [PATCH 3/3] Test width-forward derivative constants --- .../MLIR/Implementations/MathDerivatives.td | 2 +- .../batched_derivative_constant.mlir | 20 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 enzyme/test/MLIR/ForwardMode/batched_derivative_constant.mlir diff --git a/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td b/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td index 5905542eb21..aef30f672ec 100644 --- a/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td +++ b/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td @@ -22,7 +22,7 @@ def : MLIRDerivative<"math", "SqrtOp", (Op $x), >; def : MLIRDerivative<"math", "AtanOp", (Op $x), [ - (CheckedMulF (DiffeRet), (DivF (ConstantFP<"1.0","arith","ConstantOp"> $x):$one, (AddF (MulF $x, $x), $one))) + (CheckedMulF (DiffeRet), (DivF (ConstantFP<"1.0","arith","ConstantOp">):$one, (AddF (MulF $x, $x), $one))) ] >; def : MLIRDerivative<"math", "AbsFOp", (Op $x), diff --git a/enzyme/test/MLIR/ForwardMode/batched_derivative_constant.mlir b/enzyme/test/MLIR/ForwardMode/batched_derivative_constant.mlir new file mode 100644 index 00000000000..cbbb427fdb6 --- /dev/null +++ b/enzyme/test/MLIR/ForwardMode/batched_derivative_constant.mlir @@ -0,0 +1,20 @@ +// RUN: %eopt --enzyme %s | FileCheck %s + +// Regression for a width-forward derivative with an operand-less floating +// constant. The constant participates in the tangent expression and must use +// the width-stacked shadow type. +module { + func.func @atan(%x: tensor<4xf32>) -> tensor<4xf32> { + %0 = math.atan %x : tensor<4xf32> + return %0 : tensor<4xf32> + } + + func.func @main(%x: tensor<4xf32>, %dx: tensor<2x4xf32>) -> tensor<2x4xf32> { + %0 = enzyme.fwddiff @atan(%x, %dx) {activity = [#enzyme], ret_activity = [#enzyme], width = 2 : i64} : (tensor<4xf32>, tensor<2x4xf32>) -> tensor<2x4xf32> + return %0 : tensor<2x4xf32> + } +} + +// CHECK-LABEL: func.func private @fwddiffe2atan +// CHECK: arith.constant dense<1.000000e+00> : tensor<2x4xf32> +// CHECK: arith.divf