EnzymeAD / EnzymeAD/Enzyme-JAX

factor out the common matvec

Open
#1,863 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_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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.