Skip to content

Commit e359b8a

Browse files
PangorawPangoraw
authored
Bump Enzyme commit & fix complex tests (#2222)
* fix complex tests * bump enzyme commit
1 parent c801fb3 commit e359b8a

File tree

6 files changed

+22
-22
lines changed

6 files changed

+22
-22
lines changed

‎test/lit_tests/diffrules/stablehlo/abs.mlir‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs_complex outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD-COMPLEX
4-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs_complex outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE-COMPLEX
5-
// RUN: enzymexlamlir-opt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --chlo-legalize-to-stablehlo --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
1+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs_complex outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD-COMPLEX
4+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=abs_complex outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE-COMPLEX
5+
// RUN: enzymexlamlir-opt %s --enzyme --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --chlo-legalize-to-stablehlo --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
66

77
func.func @abs(%arg0 : tensor<2xf32>) -> tensor<2xf32> {
88
%0 = stablehlo.abs %arg0 : (tensor<2xf32>) -> tensor<2xf32>

‎test/lit_tests/diffrules/stablehlo/complex.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=complex outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=complex outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3-
// RUN: enzymexlamlir-opt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret
1+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=complex outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=complex outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3+
// RUN: enzymexlamlir-opt %s --enzyme --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret
44

55
func.func @complex(%a : tensor<2xf32>, %b : tensor<2xf32>) -> tensor<2xcomplex<f32>> {
66
%c = stablehlo.complex %a, %b : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xcomplex<f32>>

‎test/lit_tests/diffrules/stablehlo/imag.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=imag outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=imag outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3-
// RUN: enzymexlamlir-opt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
1+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=imag outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=imag outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3+
// RUN: enzymexlamlir-opt %s --enzyme --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
44

55
func.func @imag(%x : tensor<2xcomplex<f32>>) -> tensor<2xf32> {
66
%y = stablehlo.imag %x : (tensor<2xcomplex<f32>>) -> tensor<2xf32>

‎test/lit_tests/diffrules/stablehlo/multiply.mlir‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=REVERSE
3-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply_complex outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=FORWARD-COMPLEX
4-
// RUN enzymexlamlir-opt %s --enzyme-wrap="infn=multiply_complex outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE-COMPLEX
5-
// RUN: enzymexlamlir-opt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --arith-raise | stablehlo-translate - --interpret
2+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=REVERSE
3+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply_complex outfn= retTys=enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --arith-raise --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=FORWARD-COMPLEX
4+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=multiply_complex outfn= retTys=enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE-COMPLEX
5+
// RUN: enzymexlamlir-opt %s --enzyme --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise | stablehlo-translate - --interpret
66

77
func.func @multiply(%a : tensor<2xf32>, %b : tensor<2xf32>) -> tensor<2xf32> {
88
%c = stablehlo.multiply %a, %b : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
@@ -37,10 +37,10 @@ func.func @multiply_complex(%a : tensor<2xcomplex<f32>>, %b : tensor<2xcomplex<f
3737
// FORWARD-COMPLEX-NEXT: }
3838

3939
// REVERSE-COMPLEX: func.func @multiply_complex(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>, %arg2: tensor<2xcomplex<f32>>) -> (tensor<2xcomplex<f32>>, tensor<2xcomplex<f32>>) {
40-
// REVERSE-COMPLEX-NEXT: %0 = chlo.conj %arg2 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
41-
// REVERSE-COMPLEX-NEXT: %1 = stablehlo.multiply %0, %arg1 : tensor<2xcomplex<f32>>
40+
// REVERSE-COMPLEX-NEXT: %0 = chlo.conj %arg2 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
41+
// REVERSE-COMPLEX-NEXT: %1 = stablehlo.multiply %0, %arg1 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<2xcomplex<f32>>
4242
// REVERSE-COMPLEX-NEXT: %2 = chlo.conj %1 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
43-
// REVERSE-COMPLEX-NEXT: %3 = stablehlo.multiply %0, %arg0 : tensor<2xcomplex<f32>>
43+
// REVERSE-COMPLEX-NEXT: %3 = stablehlo.multiply %0, %arg0 {enzymexla.complex_is_purely_real = [#enzymexla<guaranteed NOTGUARANTEED>]} : tensor<2xcomplex<f32>>
4444
// REVERSE-COMPLEX-NEXT: %4 = chlo.conj %3 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
4545
// REVERSE-COMPLEX-NEXT: return %2, %4 : tensor<2xcomplex<f32>>, tensor<2xcomplex<f32>>
4646
// REVERSE-COMPLEX-NEXT: }

‎test/lit_tests/diffrules/stablehlo/real.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=real outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2-
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=real outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3-
// RUN: enzymexlamlir-opt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
1+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=real outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --arith-raise --enzyme-hlo-opt --cse | FileCheck %s --check-prefix=FORWARD
2+
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=real outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --cse --verify-each=0 | FileCheck %s --check-prefix=REVERSE
3+
// RUN: enzymexlamlir-opt %s --enzyme --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --enzyme-hlo-opt --verify-each=0 | stablehlo-translate - --interpret --allow-unregistered-dialect
44

55
func.func @real(%x : tensor<2xcomplex<f32>>) -> tensor<2xf32> {
66
%y = stablehlo.real %x : (tensor<2xcomplex<f32>>) -> tensor<2xf32>

‎workspace.bzl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
JAX_COMMIT = "8c5def0f34b37cacf27ed1114a1f958dcb2153be"
22
JAX_SHA256 = ""
33

4-
ENZYME_COMMIT = "b847f3bf85b4e74b85170781f83c36a80ce15b96"
4+
ENZYME_COMMIT = "d9db71943e66c6dff80849f01cc1d261a46a485b"
55
ENZYME_SHA256 = ""
66

77
ML_TOOLCHAIN_COMMIT = "30ef4a9096f9490e8f198faa5ce5bbddd1b72fdb"

0 commit comments

Comments
 (0)