`--enzyme-batch` produces malformed `stablehlo.constant` / `chlo.constant` (value attr not resized to batched shape)
# `--enzyme-batch` produces malformed `stablehlo.constant` / `chlo.constant` (value attr not resized to batched shape)
## Summary
When `EnzymeBatchPass` (`--enzyme-batch`) batch-clones a region that contains a
`stablehlo.constant` or `chlo.constant`, it prepends the batch dimension(s) to the op's
**result type** but copies the op's attributes verbatim. For constant ops the data lives in
the `value` attribute, so the result type and the `value` attribute end up inconsistent:
```mlir
%10 = "chlo.constant"() <{value = dense<1.1283791670955126> : tensor<f64>}> : () -> tensor<8xf64>
%11 = "chlo.constant"() <{value = dense<1.1283791670955126> : tensor<7xf64>}> : () -> tensor<8x7xf64>
```
These fail verification:
```
error: 'chlo.constant' op inferred type(s) 'tensor<f64>' are incompatible with return type(s) of operation 'tensor<8xf64>'
error: 'chlo.constant' op failed to infer returned types
```
Only `arith::ConstantOp` has a `BatchOpInterface` that resizes the attribute
(`ArithConstantOpBatchInterface`). `stablehlo::ConstantOp` and `chlo::ConstantOp` have none,
so they fall through to the generic clone path in `batchCloneBlock` and are emitted malformed.
## Root cause
`Enzyme/MLIR/Passes/EnzymeBatchPass.cpp`, `batchCloneBlock` (generic path):
```cpp
SmallVector<Type> resultTypes(src.getResultTypes().begin(),
src.getResultTypes().end());
for (auto &Ty : resultTypes) {
Ty = applyBatchSizes(Ty, batchSizes); // result type gains batch dims
}
Operation *newOp = Operation::create(
src.getLoc(), src.getName(), resultTypes, operands, src.getAttrs(), // attrs copied verbatim
mlir::PropertyRef(), successors, src.getNumRegions());
```
For a constant op, `value` is a `DenseElementsAttr` whose shaped type must equal the result
type. The result type is batched but `value` is not, producing the inconsistency above.
The generic path is taken only for ops without a `BatchOpInterface`. The single existing
constant batch rule is for arith:
`Enzyme/MLIR/Implementations/ArithAutoDiffOpInterfaceImpl.cpp`
```cpp
struct ArithConstantOpBatchInterface
: public BatchOpInterface::ExternalModel<ArithConstantOpBatchInterface, arith::ConstantOp> {
mlir::LogicalResult createBatch(Operation *src, OpBuilder &builder, IRMapping &mapper,
ArrayRef<int64_t> batchSizes) const {
SmallVector<Type> resultTypes(src->getResultTypes().begin(), src->getResultTypes().end());
for (auto &Ty : resultTypes) {
auto T = cast<TensorType>(Ty);
SmallVector<int64_t> shape(batchSizes.begin(), batchSizes.end());
shape.append(T.getShape().begin(), T.getShape().end());
Ty = T.clone(shape);
}
mlir::NamedAttrList attrs;
for (auto attr : src->getAttrs()) {
auto eattr = cast<DenseElementsAttr>(attr.getValue());
attr.setValue(eattr.resizeSplat(cast<ShapedType>(resultTypes[0]))); // <-- the missing step
attrs.append(attr);
}
auto cop = mlir::Operation::create(src->getLoc(), src->getName(), resultTypes, {},
std::move(attrs), mlir::PropertyRef(), mlir::BlockRange(), 0);
builder.insert(cop);
mapper.map(src->getResult(0), cop->getResult(0));
return success();
}
};
```
`stablehlo::ConstantOp` and `chlo::ConstantOp` need the equivalent.
## Minimal reproducer
Analogous to the existing `test/MLIR/Batch/addconst.mlir` (which passes for `arith.constant`):
```mlir
// RUN: enzymexlamlir-opt --enzyme-batch %s | FileCheck %s
module {
func.func @f(%x : tensor<3xf64>) -> tensor<3xf64> {
%cst = stablehlo.constant dense<2.1> : tensor<3xf64>
%y = stablehlo.add %x, %cst : tensor<3xf64>
return %y : tensor<3xf64>
}
func.func @df(%x : tensor<10x3xf64>) -> tensor<10x3xf64> {
%r = enzyme.batch @f(%x) { batch_shape = array<i64: 10> } : (tensor<10x3xf64>) -> (tensor<10x3xf64>)
return %r : tensor<10x3xf64>
}
}
// Expected (mirroring the arith.constant case):
// CHECK: %cst = stablehlo.constant dense<2.100000e+00> : tensor<10x3xf64>
// Actual: value attr stays tensor<3xf64> while result type becomes tensor<10x3xf64> -> verifier failure.
```
The same applies to `chlo.constant`.
## How it surfaces in practice
Reverse-mode AD of `chlo.erf` emits its derivative coefficient `erf'(x) = (2/√π)·exp(−x²)`,
with `2/√π = 1.1283791670955126` materialized as a `chlo.constant` at the per-sample shape
(`tensor<7xf64>`/`tensor<f64>`). When the surrounding function is then batched (here, an 8-sample
Monte-Carlo estimator → `batch_shape = [8]`), the constant's result type gains the leading `8`
but the `value` attr does not, yielding the malformed ops above and a hard pass-manager failure.
`stablehlo.constant`s in the same module happen to survive because later `EnzymeHLOOpt`
canonicalizations rematerialize `stablehlo::ConstantOp`s at the correct shape via `resizeSplat`;
those patterns never build `chlo::ConstantOp`, so the chlo constant from the erf decomposition is
never repaired and reaches verification malformed. Both ops are nonetheless emitted incorrectly by
the batch pass — `chlo.constant` is just the one that consistently escapes the later fixups.
## Suggested fix
Add a `BatchOpInterface` external model for `stablehlo::ConstantOp` and `chlo::ConstantOp`
mirroring `ArithConstantOpBatchInterface` (resize the `DenseElementsAttr` to the batched result
type with `resizeSplat`), registered in `StableHLOAutoDiffOpInterfaceImpl.cpp`'s
`registerInterfaces`. `resizeSplat` covers splat constants (the case here and the one the arith
rule handles); a fully general version would `broadcast_in_dim` non-splat constants along the new
leading batch dims.
Alternatively, handle constant-like attributes in the generic `batchCloneBlock` path directly, but
a per-op `BatchOpInterface` matches the existing arith precedent and is type-safe.
## Environment
- Reactant.jl (Julia), `chlo`/`stablehlo` via Enzyme-JAX.
1 条评论