EnzymeAD / EnzymeAD/Enzyme-JAX
concatenate of same value is reshape -> bcast -> reshape
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_elemmul... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<32xi64>, %arg1: tensor<48xi64>, %arg2: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
%cst = stablehlo.constant dense<2.000000e+00> : tensor<1536xf32>
%c = stablehlo.constant dense<1> : tensor<1536x2xi64>
%cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<1024x1024xf32>
%0 = stablehlo.transpose %arg2, dims = [1, 0] : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
%1 = stablehlo.concatenate %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, dim = 0 : (tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>, tensor<32xi64>) -> tensor<1536xi64>
%2 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<48xi64>) -> tensor<48x32xi64>
%3 = stablehlo.reshape %1 : (tensor<1536xi64>) -> tensor<1536x1xi64>
%4 = stablehlo.reshape %2 : (tensor<48x32xi64>) -> tensor<1536x1xi64>
%5 = stablehlo.concatenate %3, %4, dim = 1 : (tensor<1536x1xi64>, tensor<1536x1xi64>) -> tensor<1536x2xi64>
%6 = stablehlo.subtract %5, %c : tensor<1536x2xi64>
%7 = "stablehlo.scatter"(%cst_0, %6, %cst) <{scatter_dimension_numbers = #stablehlo.scatter}> ({
^bb0(%arg3: tensor, %arg4: tensor):
stablehlo.return %arg4 : tensor
}) : (tensor<1024x1024xf32>, tensor<1536x2xi64>, tensor<1536xf32>) -> tensor<1024x1024xf32>
%8 = stablehlo.multiply %7, %0 : tensor<1024x1024xf32>
%9 = stablehlo.transpose %8, dims = [1, 0] : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
return %9 : tensor<1024x1024xf32>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Research direction
No source files or tests are named. Start by running the supplied MLIR module through the project's StableHLO/MLIR path and inspect how the repeated stablehlo.concatenate is handled; done means reproducing the reported reshape-to-broadcast-to-reshape behavior and establishing the expected transformation with a regression test.
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
- 25/100