EnzymeAD / EnzymeAD/Enzyme-JAX

Slice of While(DUS) along iteration dimension

Open
#1,636 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.