EnzymeAD / EnzymeAD/Enzyme-JAX

concatenate of same value is reshape -> bcast -> reshape

Open
#1,032 1 comment 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_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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.