EnzymeAD / EnzymeAD/Enzyme-JAX

Cast `i64` indexing to lower precision automatically if possible

Open
#1,303 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

Example case (probably this is more complicated) where we can downcast

```mlir
module @reactant_findfir... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<12xi1>) -> tensor {
%c = stablehlo.constant dense<1> : tensor
%c_0 = stablehlo.constant dense<13> : tensor
%c_1 = stablehlo.constant dense : tensor
%c_2 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]> : tensor<12xi64>
%0:2 = stablehlo.reduce(%arg0 init: %c_1), (%c_2 init: %c_0) across dimensions = [0] : (tensor<12xi1>, tensor<12xi64>, tensor, tensor) -> (tensor, tensor)
reducer(%arg1: tensor, %arg3: tensor) (%arg2: tensor, %arg4: tensor) {
%2 = stablehlo.or %arg1, %arg3 : tensor
%3 = stablehlo.minimum %arg2, %arg4 : tensor
%4 = stablehlo.select %arg3, %arg4, %c_0 : tensor, tensor
%5 = stablehlo.select %arg1, %3, %4 : tensor, tensor
stablehlo.return %2, %5 : tensor, tensor
}
%1 = stablehlo.add %0#1, %c : tensor
return %1 : tensor
}
}
```

Simpler cases: (this case might be hard to prove? maybe attributes on arguments?)

```mlir
module @reactant_getindex attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<32x32xf64>, %arg1: tensor) -> tensor<32xf64> {
%c = stablehlo.constant dense<0> : tensor
%c_0 = stablehlo.constant dense<1> : tensor
%0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<32x32xf64>) -> tensor<32x32xf64>
%1 = stablehlo.subtract %arg1, %c_0 : tensor
%2 = stablehlo.convert %1 : (tensor) -> tensor
%3 = stablehlo.dynamic_slice %0, %2, %c, sizes = [1, 32] : (tensor<32x32xf64>, tensor, tensor) -> tensor<1x32xf64>
%4 = stablehlo.reshape %3 : (tensor<1x32xf64>) -> tensor<32xf64>
return %4 : tensor<32xf64>
}
}
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the two MLIR examples in the issue and compare the indexing values and types involved. Investigate how the compiler can determine that an i64 index fits a lower precision type, including the reduce example and possible argument attributes. Done means safe automatic downcasts occur where representable without changing the resulting MLIR behavior.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.