EnzymeAD / EnzymeAD/Enzyme-JAX

LU(Symmetric) \ N -> Bunch Kaufman (sytrf)

Open
#2,186 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_operation attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<3000x3000xf64> {enzymexla.memory_effects = []}, %arg1: tensor<3000x3000xf64> {enzymexla.memory_effects = []}, %arg2: tensor<1x3000xf64> {enzymexla.memory_effects = []}) -> tensor<3000x3000xf64> attributes {enzymexla.memory_effects = []} {
%cst = stablehlo.constant dense<3.3333333333333332E-4> : tensor<3000xf64>
%cst_0 = stablehlo.constant dense<3.3344448149383126E-4> : tensor<3000x3000xf64>
%cst_1 = stablehlo.constant dense<3.000000e+03> : tensor
%c = stablehlo.constant dense<1> : tensor<3000x1xi64>
%cst_2 = stablehlo.constant dense<1.000000e+00> : tensor
%cst_3 = stablehlo.constant dense<1.000000e-01> : tensor<3000xf64>
%cst_4 = stablehlo.constant dense<0.000000e+00> : tensor
%cst_5 = stablehlo.constant dense<0.000000e+00> : tensor<3000x3000xf64>
%0 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<3000x3000xf64>) -> tensor<3000x3000xf64>
%1 = stablehlo.iota dim = 0 : tensor<3000x2xi64>
%2 = "stablehlo.scatter"(%cst_5, %1, %cst_3) <{scatter_dimension_numbers = #stablehlo.scatter}> ({
^bb0(%arg3: tensor, %arg4: tensor):
stablehlo.return %arg4 : tensor
}) : (tensor<3000x3000xf64>, tensor<3000x2xi64>, tensor<3000xf64>) -> tensor<3000x3000xf64>
%3 = stablehlo.reduce(%arg0 init: %cst_4) applies stablehlo.add across dimensions = [0] : (tensor<3000x3000xf64>, tensor) -> tensor<3000xf64>
%4 = stablehlo.reduce(%arg1 init: %cst_4) applies stablehlo.add across dimensions = [0] : (tensor<3000x3000xf64>, tensor) -> tensor<3000xf64>
%5 = stablehlo.multiply %4, %cst : tensor<3000xf64>
%6 = stablehlo.reshape %5 : (tensor<3000xf64>) -> tensor<3000x1xf64>
%7 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla]} : (tensor<3000x3000xf64>, tensor<3000x3000xf64>) -> tensor<3000x3000xf64>
%8 = stablehlo.dot_general %3, %5, contracting_dims = [] x [], precision = [DEFAULT, DEFAULT] : (tensor<3000xf64>, tensor<3000xf64>) -> tensor<3000x3000xf64>
%9 = stablehlo.subtract %7, %8 {enzymexla.symmetric_matrix = [#enzymexla]} : tensor<3000x3000xf64>
%10 = enzymexla.blas.syrk %arg1, %cst_5, %cst_2, %cst_4 {output_uplo = #enzymexla.uplo, transpose = #enzymexla.transpose, uplo = #enzymexla.uplo} : (tensor<3000x3000xf64>, tensor<3000x3000xf64>, tensor, tensor) -> tensor<3000x3000xf64>
%11 = enzymexla.blas.syrk %6, %cst_5, %cst_1, %cst_4 {output_uplo = #enzymexla.uplo, uplo = #enzymexla.uplo} : (tensor<3000x1xf64>, tensor<3000x3000xf64>, tensor, tensor) -> tensor<3000x3000xf64>
%12 = stablehlo.subtract %10, %11 : tensor<3000x3000xf64>
%13 = stablehlo.multiply %12, %cst_0 : tensor<3000x3000xf64>
%14 = stablehlo.add %13, %2 : tensor<3000x3000xf64>
%15 = stablehlo.broadcast_in_dim %arg2, dims = [1, 0] : (tensor<1x3000xf64>) -> tensor<3000x3000xf64>
%16 = stablehlo.subtract %15, %0 : tensor<3000x3000xf64>
%17 = stablehlo.add %16, %2 : tensor<3000x3000xf64>
%output, %pivots, %permutation, %info = enzymexla.linalg.lu %14 : (tensor<3000x3000xf64>) -> (tensor<3000x3000xf64>, tensor<3000xi32>, tensor<3000xi32>, tensor)
%18 = stablehlo.convert %permutation : (tensor<3000xi32>) -> tensor<3000xi64>
%19 = stablehlo.reshape %18 : (tensor<3000xi64>) -> tensor<3000x1xi64>
%20 = stablehlo.subtract %19, %c : tensor<3000x1xi64>
%21 = "stablehlo.gather"(%17, %20) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<3000x3000xf64>, tensor<3000x1xi64>) -> tensor<3000x3000xf64>
%22 = "stablehlo.triangular_solve"(%output, %21) <{left_side = true, lower = true, transpose_a = #stablehlo, unit_diagonal = true}> : (tensor<3000x3000xf64>, tensor<3000x3000xf64>) -> tensor<3000x3000xf64>
%23 = "stablehlo.triangular_solve"(%output, %22) <{left_side = true, lower = false, transpose_a = #stablehlo, unit_diagonal = false}> {enzymexla.symmetric_matrix = [#enzymexla]} : (tensor<3000x3000xf64>, tensor<3000x3000xf64>) -> tensor<3000x3000xf64>
%24 = stablehlo.dot_general %23, %9, contracting_dims = [0] x [1], precision = [DEFAULT, DEFAULT] : (tensor<3000x3000xf64>, tensor<3000x3000xf64>) -> tensor<3000x3000xf64>
%25 = stablehlo.multiply %cst_0, %24 : tensor<3000x3000xf64>
%26 = stablehlo.add %arg0, %25 {enzymexla.symmetric_matrix = [#enzymexla]} : tensor<3000x3000xf64>
return %26 : tensor<3000x3000xf64>
}
}
```

from https://discourse.julialang.org/t/gpu-performance-degradation-due-to-a-b-calculation/136072/8

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.