EnzymeAD / EnzymeAD/Enzyme-JAX

`PadDotGeneral` is busted

Open
#1,445 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

```mlir
// ./bazel-bin/enzymexlamlir-opt --pass-pipeline="any(enzyme-hlo-generate-td{patterns=pad_dot_general<1>(0);pad_dot_general<1>(1)},transform-interpreter,enzyme-hlo-remove-transform)" test/lit_tests/dus_to_dynamic_pad.mlir
module {
func.func @fn(%arg0: tensor {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) -> (tensor<32x32xf32>, tensor) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
%c = stablehlo.constant dense<3> : tensor<32x32xi64>
%c_0 = stablehlo.constant dense<2> : tensor<32x32xi64>
%cst = stablehlo.constant dense<0.000000e+00> : tensor
%0 = stablehlo.reshape %arg0 : (tensor) -> tensor<1x1xf32>
%1 = stablehlo.pad %0, %cst, low = [2, 3], high = [29, 28], interior = [0, 0] : (tensor<1x1xf32>, tensor) -> tensor<32x32xf32>
%2 = stablehlo.convert %c_0 : (tensor<32x32xi64>) -> tensor<32x32xf32>
%3 = stablehlo.multiply %2, %1 : tensor<32x32xf32>
%4 = stablehlo.convert %c : (tensor<32x32xi64>) -> tensor<32x32xf32>
%5 = stablehlo.subtract %3, %4 : tensor<32x32xf32>
%6 = stablehlo.pad %0, %cst, low = [3, 2], high = [28, 29], interior = [0, 0] : (tensor<1x1xf32>, tensor) -> tensor<32x32xf32>
%7 = stablehlo.dot_general %6, %5, contracting_dims = [0] x [1], precision = [DEFAULT, DEFAULT] : (tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32>
return %7, %arg0 : tensor<32x32xf32>, tensor
}
}
```

```mlir
test/lit_tests/dus_to_dynamic_pad.mlir:14:10: error: contracting dimension sizes must match for lhs/rhs
%7 = stablehlo.dot_general %6, %5, contracting_dims = [0] x [1], precision = [DEFAULT, DEFAULT] : (tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32>
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Run the bazel-bin/enzymexlamlir-opt command from the issue against test/lit_tests/dus_to_dynamic_pad.mlir and inspect the PadDotGeneral path used by the pad_dot_general patterns. Trace the reported contracting-dimension mismatch at the stablehlo.dot_general operation; done means the reproducer no longer emits that diagnostic and the relevant lit test passes.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.