EnzymeAD / EnzymeAD/Enzyme-JAX
factor out the common matvec
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_gesummv attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor {enzymexla.memory_effects = []}, %arg1: tensor {enzymexla.memory_effects = []}, %arg2: tensor<64x64xf32> {enzymexla.memory_effects = []}, %arg3: tensor<64x64xf32> {enzymexla.memory_effects = []}, %arg4: tensor<64xf32> {enzymexla.memory_effects = []}) -> tensor<64xf32> attributes {enzymexla.memory_effects = []} {
%0 = stablehlo.dot_general %arg3, %arg4, contracting_dims = [0] x [0] : (tensor<64x64xf32>, tensor<64xf32>) -> tensor<64xf32>
%1 = stablehlo.dot_general %arg2, %arg4, contracting_dims = [0] x [0] : (tensor<64x64xf32>, tensor<64xf32>) -> tensor<64xf32>
%2 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor) -> tensor<64xf32>
%3 = stablehlo.multiply %2, %1 : tensor<64xf32>
%4 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor) -> tensor<64xf32>
%5 = stablehlo.multiply %4, %0 : tensor<64xf32>
%6 = stablehlo.add %3, %5 : tensor<64xf32>
return %6 : tensor<64xf32>
}
}
```
this is computing $\alpha A x + \beta B x$. we should convert this into a single matmul
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.