EnzymeAD / EnzymeAD/Enzyme-JAX
[polybench] blas trmm raising failure
- Dominant language
- MLIR
- Stars
- 133
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_kernel_... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor {enzymexla.memory_effects = []}, %arg1: tensor<64x64xf32> {enzymexla.memory_effects = []}, %arg2: tensor<32x64xf32> {enzymexla.memory_effects = [], tf.aliasing_output = 0 : i32}) -> tensor<32x64xf32> attributes {enzymexla.memory_effects = []} {
%c = stablehlo.constant dense<1> : tensor<64x64xi64>
%c_0 = stablehlo.constant dense<32> : tensor
%c_1 = stablehlo.constant dense<64> : tensor
%c_2 = stablehlo.constant dense<1> : tensor
%c_3 = stablehlo.constant dense<0> : tensor
%c_4 = stablehlo.constant dense<1> : tensor
%0 = stablehlo.reshape %arg0 : (tensor) -> tensor<1x1xf32>
%1 = stablehlo.iota dim = 1 : tensor<64x64xi64>
%2 = stablehlo.add %c, %1 : tensor<64x64xi64>
%3 = stablehlo.iota dim = 0 : tensor<64x64xi64>
%4 = stablehlo.add %c, %3 : tensor<64x64xi64>
%5 = stablehlo.compare GT, %4, %2 : (tensor<64x64xi64>, tensor<64x64xi64>) -> tensor<64x64xi1>
%6 = stablehlo.convert %5 : (tensor<64x64xi1>) -> tensor<64x64xf32>
%7 = stablehlo.multiply %6, %arg1 : tensor<64x64xf32>
%8:2 = stablehlo.while(%iterArg = %c_3, %iterArg_5 = %arg2) : tensor, tensor<32x64xf32> attributes {enzyme.disable_mincut, enzymexla.symmetric_matrix = [#enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla]}
cond {
%9 = stablehlo.compare LT, %iterArg, %c_1 : (tensor, tensor) -> tensor
stablehlo.return %9 : tensor
} do {
%9 = stablehlo.add %c_4, %iterArg {enzymexla.bounds = [[1, 64]]} : tensor
%10 = stablehlo.convert %9 {enzymexla.bounds = [[1, 64]]} : (tensor) -> tensor
%11 = stablehlo.subtract %10, %c_2 {enzymexla.bounds = [[0, 63]]} : tensor
%12:2 = stablehlo.while(%iterArg_6 = %c_3, %iterArg_7 = %iterArg_5) : tensor, tensor<32x64xf32> attributes {enzyme.disable_mincut}
cond {
%13 = stablehlo.compare LT, %iterArg_6, %c_0 : (tensor, tensor) -> tensor
stablehlo.return %13 : tensor
} do {
%13 = stablehlo.transpose %iterArg_7, dims = [1, 0] : (tensor<32x64xf32>) -> tensor<64x32xf32>
%14 = stablehlo.add %c_4, %iterArg_6 {enzymexla.bounds = [[1, 32]]} : tensor
%15 = stablehlo.convert %14 {enzymexla.bounds = [[1, 32]]} : (tensor) -> tensor
%16 = stablehlo.subtract %15, %c_2 {enzymexla.bounds = [[0, 31]]} : tensor
%17 = stablehlo.dynamic_slice %iterArg_5, %16, %11, sizes = [1, 1] : (tensor<32x64xf32>, tensor, tensor) -> tensor<1x1xf32>
%18:2 = stablehlo.while(%iterArg_8 = %c_3, %iterArg_9 = %17) : tensor, tensor<1x1xf32> attributes {enzyme.disable_mincut}
cond {
%21 = stablehlo.compare LT, %iterArg_8, %c_1 : (tensor, tensor) -> tensor
stablehlo.return %21 : tensor
} do {
%21 = stablehlo.dynamic_update_slice %13, %iterArg_9, %11, %16 : (tensor<64x32xf32>, tensor<1x1xf32>, tensor, tensor) -> tensor<64x32xf32>
%22 = stablehlo.add %c_4, %iterArg_8 {enzymexla.bounds = [[1, 64]]} : tensor
%23 = stablehlo.convert %22 {enzymexla.bounds = [[1, 64]]} : (tensor) -> tensor
%24 = stablehlo.subtract %23, %c_2 {enzymexla.bounds = [[0, 63]]} : tensor
%25 = stablehlo.dynamic_slice %21, %24, %16, sizes = [1, 1] : (tensor<64x32xf32>, tensor, tensor) -> tensor<1x1xf32>
%26 = stablehlo.dynamic_slice %7, %iterArg, %iterArg_8, sizes = [1, 1] : (tensor<64x64xf32>, tensor, tensor) -> tensor<1x1xf32>
%27 = stablehlo.multiply %26, %25 : tensor<1x1xf32>
%28 = stablehlo.add %iterArg_9, %27 : tensor<1x1xf32>
stablehlo.return %22, %28 : tensor, tensor<1x1xf32>
}
%19 = stablehlo.multiply %0, %18#1 : tensor<1x1xf32>
%20 = stablehlo.dynamic_update_slice %iterArg_7, %19, %16, %11 : (tensor<32x64xf32>, tensor<1x1xf32>, tensor, tensor) -> tensor<32x64xf32>
stablehlo.return %14, %20 : tensor, tensor<32x64xf32>
}
stablehlo.return %9, %12#1 : tensor, tensor<32x64xf32>
}
return %8#1 : tensor<32x64xf32>
}
}
```
```
%21 = stablehlo.dynamic_update_slice %13, %iterArg_9, %11, %16 : (tensor<64x32xf32>, tensor<1x1xf32>, tensor, tensor) -> tensor<64x32xf32>
%25 = stablehlo.dynamic_slice %21, %24, %16, sizes = [1, 1] : (tensor<64x32xf32>, tensor, tensor) -> tensor<1x1xf32>
```
we can rewrite this as `select(%24 == %11, %iterArg_9, dynamic_slice(%13, %24, %16))`
Contributor guide
No contributing guide indexed for this repository
Research direction
Start from the attached MLIR for the polybench BLAS TRMM failure, focusing on the dynamic_update_slice and dynamic_slice sequence shown near %21 and %25. Compare that behavior with the proposed select rewrite; done means the TRMM case no longer raises the reported failure.
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
- Mostly clear
- Newbie friendliness
- 35/100