Skip to content

ProbProg: numerical fix #2915

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sbrantq merged 1 commit into from
Jul 2, 2026
Merged

ProbProg: numerical fix #2915

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions enzyme/Enzyme/MLIR/Interfaces/HMCUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1622,20 +1622,31 @@ impulse::updateCheckpoints(OpBuilder &builder, Location loc, Value leafIdx,
}

DualAveragingState impulse::initDualAveraging(OpBuilder &builder, Location loc,
Value stepSize) {
Value stepSize,
bool logOfProduct) {
auto stepSizeType = cast<RankedTensorType>(stepSize.getType());
auto elemType = stepSizeType.getElementType();
auto scalarType = RankedTensorType::get({}, elemType);
auto i64TensorType = RankedTensorType::get({}, builder.getI64Type());

// prox_center = log(10) + log(step_size)
auto log10Const = arith::ConstantOp::create(
builder, loc, scalarType,
DenseElementsAttr::get(scalarType,
builder.getFloatAttr(elemType, std::log(10.0))));
Value logStepSize = math::LogOp::create(builder, loc, stepSize);
Value proxCenter =
arith::AddFOp::create(builder, loc, log10Const, logStepSize);
// prox_center: log(10 * step_size) at warmup start, or log(10) +
// log(step_size) at window-end re-inits
Value proxCenter;
if (logOfProduct) {
auto tenConst = arith::ConstantOp::create(
builder, loc, scalarType,
DenseElementsAttr::get(scalarType,
builder.getFloatAttr(elemType, 10.0)));
Value scaled = arith::MulFOp::create(builder, loc, tenConst, stepSize);
proxCenter = math::LogOp::create(builder, loc, scaled);
} else {
auto log10Const = arith::ConstantOp::create(
builder, loc, scalarType,
DenseElementsAttr::get(scalarType,
builder.getFloatAttr(elemType, std::log(10.0))));
Value logStepSize = math::LogOp::create(builder, loc, stepSize);
proxCenter = arith::AddFOp::create(builder, loc, log10Const, logStepSize);
}

auto zeroConst = arith::ConstantOp::create(
builder, loc, scalarType,
Expand Down
2 changes: 1 addition & 1 deletion enzyme/Enzyme/MLIR/Interfaces/HMCUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,7 +386,7 @@ struct DualAveragingConfig {

/// Initialize dual averaging state from initial step size.
DualAveragingState initDualAveraging(OpBuilder &builder, Location loc,
Value stepSize);
Value stepSize, bool logOfProduct = false);

/// Update dual averaging state with observed acceptance probability.
DualAveragingState updateDualAveraging(OpBuilder &builder, Location loc,
Expand Down
6 changes: 4 additions & 2 deletions enzyme/Enzyme/MLIR/Passes/ExpandImpulsePass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -909,7 +909,8 @@ struct ExpandImpulsePass
hasAdaptationIn
? DualAveragingState::fromValues(
ArrayRef<Value>(adaptationInVals).take_front(5))
: initDualAveraging(rewriter, loc, initialStepSize);
: initDualAveraging(rewriter, loc, initialStepSize,
/*logOfProduct=*/true);

WelfordState welfordState;
WelfordConfig welfordConfig;
Expand Down Expand Up @@ -1260,7 +1261,8 @@ struct ExpandImpulsePass
hasAdaptationIn
? DualAveragingState::fromValues(
ArrayRef<Value>(adaptationInVals).take_front(5))
: initDualAveraging(rewriter, loc, stepSize);
: initDualAveraging(rewriter, loc, stepSize,
/*logOfProduct=*/true);
WelfordState w =
hasAdaptationIn
? WelfordState::fromValues(
Expand Down
14 changes: 6 additions & 8 deletions enzyme/test/MLIR/Impulse/mcmc_warmup.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,6 @@ module {
// CHECK-DAG: %[[HALF:.+]] = arith.constant dense<5.000000e-01> : tensor<f64>
// CHECK-DAG: %[[ZERO_1D:.+]] = arith.constant dense<0.000000e+00> : tensor<1xf64>
// CHECK-DAG: %[[ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
// CHECK-DAG: %[[EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
// CHECK-DAG: %[[ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1xf64>
// CHECK-DAG: %[[C9:.+]] = arith.constant dense<9> : tensor<i64>
// CHECK-DAG: %[[C10:.+]] = arith.constant dense<10> : tensor<i64>
Expand Down Expand Up @@ -116,7 +115,7 @@ module {
// CHECK-NEXT: } attributes {activity = [#enzyme<activity enzyme_active>], ret_activity = [#enzyme<activity enzyme_active>, #enzyme<activity enzyme_const>]}
//
// --- Warmup loop: 16 iter_args ---
// CHECK-NEXT: %[[WARMUP:.+]]:16 = impulse.for(%[[C0]] : tensor<i64>) to(%[[C10]] : tensor<i64>) step(%[[C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[ZERO_F]], %[[ZERO_F]], %[[ZERO_F]], %[[C0]], %[[EPS_INIT]], %[[ZERO_1D]], %[[ZERO_1D]], %[[C0]], %[[C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64> {
// CHECK-NEXT: %[[WARMUP:.+]]:16 = impulse.for(%[[C0]] : tensor<i64>) to(%[[C10]] : tensor<i64>) step(%[[C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[ZERO_F]], %[[ZERO_F]], %[[ZERO_F]], %[[C0]], %[[ZERO_F]], %[[ZERO_1D]], %[[ZERO_1D]], %[[C0]], %[[C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64> {
// CHECK-NEXT: ^bb0(%[[WI:.+]]: tensor<i64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<2xui64>, %{{.+}}: tensor<f64>, %[[INV_MASS:.+]]: tensor<1x1xf64>, %[[MASS_SQRT:.+]]: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<1xf64>, %{{.+}}: tensor<1xf64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<i64>):
//
// --- Momentum sampling with mass matrix ---
Expand Down Expand Up @@ -275,7 +274,6 @@ module {
// CHECK-DAG: %[[S_GAMMA:.+]] = arith.constant dense<5.000000e-02> : tensor<f64>
// CHECK-DAG: %[[S_TARGET:.+]] = arith.constant dense<8.000000e-01> : tensor<f64>
// CHECK-DAG: %[[S_T0:.+]] = arith.constant dense<1.000000e+01> : tensor<f64>
// CHECK-DAG: %[[S_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
// CHECK-DAG: %[[S_C10:.+]] = arith.constant dense<10> : tensor<i64>
// CHECK-DAG: %[[S_C9:.+]] = arith.constant dense<9> : tensor<i64>
// CHECK-DAG: %[[S_C1:.+]] = arith.constant dense<1> : tensor<i64>
Expand All @@ -292,7 +290,7 @@ module {
// CHECK: }
//
// --- Warmup loop: 13 iter_args (no Welford mean/m2/n) ---
// CHECK: %[[S_WARMUP:.+]]:13 = impulse.for(%[[S_C0]] : tensor<i64>) to(%[[S_C10]] : tensor<i64>) step(%[[S_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[S_ONE_1D]], %[[S_ONE_1D]], %[[S_ZERO_F]], %[[S_ZERO_F]], %[[S_ZERO_F]], %[[S_C0]], %[[S_EPS_INIT]], %[[S_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64> {
// CHECK: %[[S_WARMUP:.+]]:13 = impulse.for(%[[S_C0]] : tensor<i64>) to(%[[S_C10]] : tensor<i64>) step(%[[S_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[S_ONE_1D]], %[[S_ONE_1D]], %[[S_ZERO_F]], %[[S_ZERO_F]], %[[S_ZERO_F]], %[[S_C0]], %[[S_ZERO_F]], %[[S_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64> {
// CHECK-NEXT: ^bb0(%[[S_WI:.+]]: tensor<i64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<2xui64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>):
//
// --- Momentum sampling ---
Expand Down Expand Up @@ -386,7 +384,7 @@ module {
// CHECK-DAG: %[[M_C0:.+]] = arith.constant dense<0> : tensor<i64>
// CHECK-DAG: %[[M_ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1xf64>
// CHECK-DAG: %[[M_ZERO_1D:.+]] = arith.constant dense<0.000000e+00> : tensor<1xf64>
// CHECK-DAG: %[[M_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
// CHECK-DAG: %[[M_ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
//
// --- Init ---
// CHECK: impulse.randomSplit
Expand All @@ -395,7 +393,7 @@ module {
// CHECK: } attributes
//
// --- Warmup loop: 16 iter_args (includes Welford state) ---
// CHECK-NEXT: %[[M_WARMUP:.+]]:16 = impulse.for(%[[M_C0]] : tensor<i64>) to(%[[M_C10]] : tensor<i64>) step(%[[M_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[M_C0]], %[[M_EPS_INIT]], %[[M_ZERO_1D]], %[[M_ZERO_1D]], %[[M_C0]], %[[M_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64> {
// CHECK-NEXT: %[[M_WARMUP:.+]]:16 = impulse.for(%[[M_C0]] : tensor<i64>) to(%[[M_C10]] : tensor<i64>) step(%[[M_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[M_C0]], %[[M_ZERO_F]], %[[M_ZERO_1D]], %[[M_ZERO_1D]], %[[M_C0]], %[[M_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<1xf64>, tensor<1xf64>, tensor<i64>, tensor<i64> {
// CHECK-NEXT: ^bb0(%[[M_WI:.+]]: tensor<i64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<2xui64>, %{{.+}}: tensor<f64>, %[[M_INV:.+]]: tensor<1x1xf64>, %[[M_SQRT:.+]]: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<1xf64>, %{{.+}}: tensor<1xf64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<i64>):
//
// --- Momentum sampling with mass matrix ---
Expand Down Expand Up @@ -481,7 +479,7 @@ module {
// CHECK-DAG: %[[N_C9:.+]] = arith.constant dense<9> : tensor<i64>
// CHECK-DAG: %[[N_C1:.+]] = arith.constant dense<1> : tensor<i64>
// CHECK-DAG: %[[N_C0:.+]] = arith.constant dense<0> : tensor<i64>
// CHECK-DAG: %[[N_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
// CHECK-DAG: %[[N_ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
// CHECK-DAG: %[[N_ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1x1xf64>
//
// --- Init ---
Expand All @@ -491,7 +489,7 @@ module {
// CHECK: } attributes
//
// --- Warmup loop: 13 iter_args ---
// CHECK-NEXT: %[[N_WARMUP:.+]]:13 = impulse.for(%[[N_C0]] : tensor<i64>) to(%[[N_C10]] : tensor<i64>) step(%[[N_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[N_ONE_1D]], %[[N_ONE_1D]], %{{.+}}, %{{.+}}, %{{.+}}, %[[N_C0]], %[[N_EPS_INIT]], %[[N_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64> {
// CHECK-NEXT: %[[N_WARMUP:.+]]:13 = impulse.for(%[[N_C0]] : tensor<i64>) to(%[[N_C10]] : tensor<i64>) step(%[[N_C1]] : tensor<i64>) iter_args(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %{{.+}}, %[[N_ONE_1D]], %[[N_ONE_1D]], %{{.+}}, %{{.+}}, %{{.+}}, %[[N_C0]], %[[N_ZERO_F]], %[[N_C0]] : tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64>) -> tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<2xui64>, tensor<f64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<f64>, tensor<f64>, tensor<f64>, tensor<i64>, tensor<f64>, tensor<i64> {
// CHECK-NEXT: ^bb0(%[[N_WI:.+]]: tensor<i64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<2xui64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<1x1xf64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>, %{{.+}}: tensor<f64>, %{{.+}}: tensor<i64>):
//
// --- Momentum sampling ---
Expand Down
Loading