EnzymeAD / EnzymeAD/Enzyme-JAX
Slice of While(DUS) along iteration dimension
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_Recurre... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<12x16x4xf32> {enzymexla.memory_effects = []}, %arg1: tensor<4x4xf32> {enzymexla.memory_effects = []}, %arg2: tensor<4x4xf32> {enzymexla.memory_effects = []}, %arg3: tensor<4xf32> {enzymexla.memory_effects = []}, %arg4: tensor<4xf32> {enzymexla.memory_effects = []}, %arg5: tensor<2xui64> {enzymexla.memory_effects = []}) -> (tensor<12x4xf32>, tensor<2xui64>) attributes {enzymexla.memory_effects = []} {
%c = stablehlo.constant dense<0> : tensor
%c_0 = stablehlo.constant dense<16> : tensor
%cst = stablehlo.constant dense<0.000000e+00> : tensor<16x12x4xf32>
%cst_1 = stablehlo.constant dense<0.000000e+00> : tensor<4x12xf32>
%c_2 = stablehlo.constant dense<1> : tensor
%c_3 = stablehlo.constant dense<1> : tensor
%c_4 = stablehlo.constant dense<0> : tensor
%0 = stablehlo.broadcast_in_dim %arg4, dims = [0] : (tensor<4xf32>) -> tensor<4x12xf32>
%1 = stablehlo.broadcast_in_dim %arg1, dims = [1, 2] : (tensor<4x4xf32>) -> tensor<16x4x4xf32>
%2 = stablehlo.dot_general %1, %arg0, batching_dims = [0] x [1], contracting_dims = [1] x [2], precision = [DEFAULT, DEFAULT] : (tensor<16x4x4xf32>, tensor<12x16x4xf32>) -> tensor<16x4x12xf32>
%3 = stablehlo.broadcast_in_dim %arg3, dims = [1] : (tensor<4xf32>) -> tensor<16x4x12xf32>
%4 = stablehlo.add %2, %3 : tensor<16x4x12xf32>
%5:3 = stablehlo.while(%iterArg = %c_4, %iterArg_5 = %cst, %iterArg_6 = %cst_1) : tensor, tensor<16x12x4xf32>, tensor<4x12xf32> attributes {enzyme.disable_mincut, enzymexla.symmetric_matrix = [#enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla, #enzymexla]}
cond {
%8 = stablehlo.compare LT, %iterArg, %c_0 : (tensor, tensor) -> tensor
stablehlo.return %8 : tensor
} do {
%8 = stablehlo.add %c_3, %iterArg : tensor
%9 = stablehlo.convert %8 : (tensor) -> tensor
%10 = stablehlo.subtract %9, %c_2 : tensor
%11 = stablehlo.dot_general %arg2, %iterArg_6, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<4x4xf32>, tensor<4x12xf32>) -> tensor<4x12xf32>
%12 = stablehlo.add %11, %0 : tensor<4x12xf32>
%13 = stablehlo.dynamic_slice %4, %iterArg, %c_4, %c_4, sizes = [1, 4, 12] : (tensor<16x4x12xf32>, tensor, tensor, tensor) -> tensor<1x4x12xf32>
%14 = stablehlo.reshape %13 : (tensor<1x4x12xf32>) -> tensor<4x12xf32>
%15 = stablehlo.add %12, %14 : tensor<4x12xf32>
%16 = stablehlo.tanh %15 : tensor<4x12xf32>
%17 = stablehlo.broadcast_in_dim %16, dims = [2, 1] : (tensor<4x12xf32>) -> tensor<1x12x4xf32>
%18 = stablehlo.dynamic_update_slice %iterArg_5, %17, %10, %c, %c : (tensor<16x12x4xf32>, tensor<1x12x4xf32>, tensor, tensor, tensor) -> tensor<16x12x4xf32>
stablehlo.return %8, %18, %16 : tensor, tensor<16x12x4xf32>, tensor<4x12xf32>
}
%6 = stablehlo.slice %5#1 [15:16, 0:12, 0:4] : (tensor<16x12x4xf32>) -> tensor<1x12x4xf32>
%7 = stablehlo.reshape %6 : (tensor<1x12x4xf32>) -> tensor<12x4xf32>
return %7, %arg5 : tensor<12x4xf32>, tensor<2xui64>
}
}
```
We don't need to cache the full `tensor<16x12x4xf32>` if only a specific iteration value is being consumed later.
Contributor guide
No contributing guide indexed for this repository
Research direction
The issue shows a StableHLO/MLIR while loop using dynamic_update_slice and a final slice, but names no source file or test. Start by tracing the transformation handling this loop and its consumed iteration value; done means avoiding the full tensor cache while preserving the shown output.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100