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 diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp index b7e3ed27407..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(), @@ -618,7 +619,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 to preserve the + // tangent shape. + shadowType = true; } else { if (resultRoot->getArgName(0)) { auto name = resultRoot->getArgName(0)->getAsUnquotedString(); @@ -665,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(), @@ -695,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");