Skip to content

Commit 84f56fc

Browse files
sbrantqsbrantq
authored
fix (#2915)
1 parent cc724cb commit 84f56fc

4 files changed

Lines changed: 31 additions & 20 deletions

File tree

‎enzyme/Enzyme/MLIR/Interfaces/HMCUtils.cpp‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1622,20 +1622,31 @@ impulse::updateCheckpoints(OpBuilder &builder, Location loc, Value leafIdx,
16221622
}
16231623

16241624
DualAveragingState impulse::initDualAveraging(OpBuilder &builder, Location loc,
1625-
Value stepSize) {
1625+
Value stepSize,
1626+
bool logOfProduct) {
16261627
auto stepSizeType = cast<RankedTensorType>(stepSize.getType());
16271628
auto elemType = stepSizeType.getElementType();
16281629
auto scalarType = RankedTensorType::get({}, elemType);
16291630
auto i64TensorType = RankedTensorType::get({}, builder.getI64Type());
16301631

1631-
// prox_center = log(10) + log(step_size)
1632-
auto log10Const = arith::ConstantOp::create(
1633-
builder, loc, scalarType,
1634-
DenseElementsAttr::get(scalarType,
1635-
builder.getFloatAttr(elemType, std::log(10.0))));
1636-
Value logStepSize = math::LogOp::create(builder, loc, stepSize);
1637-
Value proxCenter =
1638-
arith::AddFOp::create(builder, loc, log10Const, logStepSize);
1632+
// prox_center: log(10 * step_size) at warmup start, or log(10) +
1633+
// log(step_size) at window-end re-inits
1634+
Value proxCenter;
1635+
if (logOfProduct) {
1636+
auto tenConst = arith::ConstantOp::create(
1637+
builder, loc, scalarType,
1638+
DenseElementsAttr::get(scalarType,
1639+
builder.getFloatAttr(elemType, 10.0)));
1640+
Value scaled = arith::MulFOp::create(builder, loc, tenConst, stepSize);
1641+
proxCenter = math::LogOp::create(builder, loc, scaled);
1642+
} else {
1643+
auto log10Const = arith::ConstantOp::create(
1644+
builder, loc, scalarType,
1645+
DenseElementsAttr::get(scalarType,
1646+
builder.getFloatAttr(elemType, std::log(10.0))));
1647+
Value logStepSize = math::LogOp::create(builder, loc, stepSize);
1648+
proxCenter = arith::AddFOp::create(builder, loc, log10Const, logStepSize);
1649+
}
16391650

16401651
auto zeroConst = arith::ConstantOp::create(
16411652
builder, loc, scalarType,

‎enzyme/Enzyme/MLIR/Interfaces/HMCUtils.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -386,7 +386,7 @@ struct DualAveragingConfig {
386386

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

391391
/// Update dual averaging state with observed acceptance probability.
392392
DualAveragingState updateDualAveraging(OpBuilder &builder, Location loc,

‎enzyme/Enzyme/MLIR/Passes/ExpandImpulsePass.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,8 @@ struct ExpandImpulsePass
909909
hasAdaptationIn
910910
? DualAveragingState::fromValues(
911911
ArrayRef<Value>(adaptationInVals).take_front(5))
912-
: initDualAveraging(rewriter, loc, initialStepSize);
912+
: initDualAveraging(rewriter, loc, initialStepSize,
913+
/*logOfProduct=*/true);
913914

914915
WelfordState welfordState;
915916
WelfordConfig welfordConfig;
@@ -1260,7 +1261,8 @@ struct ExpandImpulsePass
12601261
hasAdaptationIn
12611262
? DualAveragingState::fromValues(
12621263
ArrayRef<Value>(adaptationInVals).take_front(5))
1263-
: initDualAveraging(rewriter, loc, stepSize);
1264+
: initDualAveraging(rewriter, loc, stepSize,
1265+
/*logOfProduct=*/true);
12641266
WelfordState w =
12651267
hasAdaptationIn
12661268
? WelfordState::fromValues(

‎enzyme/test/MLIR/Impulse/mcmc_warmup.mlir‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,6 @@ module {
8282
// CHECK-DAG: %[[HALF:.+]] = arith.constant dense<5.000000e-01> : tensor<f64>
8383
// CHECK-DAG: %[[ZERO_1D:.+]] = arith.constant dense<0.000000e+00> : tensor<1xf64>
8484
// CHECK-DAG: %[[ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
85-
// CHECK-DAG: %[[EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
8685
// CHECK-DAG: %[[ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1xf64>
8786
// CHECK-DAG: %[[C9:.+]] = arith.constant dense<9> : tensor<i64>
8887
// CHECK-DAG: %[[C10:.+]] = arith.constant dense<10> : tensor<i64>
@@ -116,7 +115,7 @@ module {
116115
// CHECK-NEXT: } attributes {activity = [#enzyme<activity enzyme_active>], ret_activity = [#enzyme<activity enzyme_active>, #enzyme<activity enzyme_const>]}
117116
//
118117
// --- Warmup loop: 16 iter_args ---
119-
// 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> {
118+
// 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> {
120119
// 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>):
121120
//
122121
// --- Momentum sampling with mass matrix ---
@@ -275,7 +274,6 @@ module {
275274
// CHECK-DAG: %[[S_GAMMA:.+]] = arith.constant dense<5.000000e-02> : tensor<f64>
276275
// CHECK-DAG: %[[S_TARGET:.+]] = arith.constant dense<8.000000e-01> : tensor<f64>
277276
// CHECK-DAG: %[[S_T0:.+]] = arith.constant dense<1.000000e+01> : tensor<f64>
278-
// CHECK-DAG: %[[S_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
279277
// CHECK-DAG: %[[S_C10:.+]] = arith.constant dense<10> : tensor<i64>
280278
// CHECK-DAG: %[[S_C9:.+]] = arith.constant dense<9> : tensor<i64>
281279
// CHECK-DAG: %[[S_C1:.+]] = arith.constant dense<1> : tensor<i64>
@@ -292,7 +290,7 @@ module {
292290
// CHECK: }
293291
//
294292
// --- Warmup loop: 13 iter_args (no Welford mean/m2/n) ---
295-
// 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> {
293+
// 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> {
296294
// 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>):
297295
//
298296
// --- Momentum sampling ---
@@ -386,7 +384,7 @@ module {
386384
// CHECK-DAG: %[[M_C0:.+]] = arith.constant dense<0> : tensor<i64>
387385
// CHECK-DAG: %[[M_ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1xf64>
388386
// CHECK-DAG: %[[M_ZERO_1D:.+]] = arith.constant dense<0.000000e+00> : tensor<1xf64>
389-
// CHECK-DAG: %[[M_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
387+
// CHECK-DAG: %[[M_ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
390388
//
391389
// --- Init ---
392390
// CHECK: impulse.randomSplit
@@ -395,7 +393,7 @@ module {
395393
// CHECK: } attributes
396394
//
397395
// --- Warmup loop: 16 iter_args (includes Welford state) ---
398-
// 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> {
396+
// 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> {
399397
// 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>):
400398
//
401399
// --- Momentum sampling with mass matrix ---
@@ -481,7 +479,7 @@ module {
481479
// CHECK-DAG: %[[N_C9:.+]] = arith.constant dense<9> : tensor<i64>
482480
// CHECK-DAG: %[[N_C1:.+]] = arith.constant dense<1> : tensor<i64>
483481
// CHECK-DAG: %[[N_C0:.+]] = arith.constant dense<0> : tensor<i64>
484-
// CHECK-DAG: %[[N_EPS_INIT:.+]] = arith.constant dense<4.44{{.+}}> : tensor<f64>
482+
// CHECK-DAG: %[[N_ZERO_F:.+]] = arith.constant dense<0.000000e+00> : tensor<f64>
485483
// CHECK-DAG: %[[N_ONE_1D:.+]] = arith.constant dense<1.000000e+00> : tensor<1x1xf64>
486484
//
487485
// --- Init ---
@@ -491,7 +489,7 @@ module {
491489
// CHECK: } attributes
492490
//
493491
// --- Warmup loop: 13 iter_args ---
494-
// 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> {
492+
// 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> {
495493
// 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>):
496494
//
497495
// --- Momentum sampling ---

0 commit comments

Comments
 (0)