diffsl 0.12.4

A compiler for a domain-specific language for ordinary differential equations (ODE).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

// RUN: %eopt --enzyme %s | FileCheck %s

module {
  func.func @nrm2(%arg0: memref<?xf64>) -> f64 {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %cst = arith.constant 0.000000e+00 : f64
    %dim = memref.dim %arg0, %c0 : memref<?xf64>
    %1 = scf.parallel (%arg1) = (%c0) to (%dim) step (%c1) init (%cst) -> (f64) {
      %2 = memref.load %arg0[%arg1] : memref<?xf64>
      %3 = arith.mulf %2, %2 : f64
      scf.reduce(%3 : f64) {
        ^bb0(%arg9: f64, %arg10: f64):
          %9 = arith.addf %arg9, %arg10 : f64
          scf.reduce.return %9 : f64
      }
    }
    return %1 : f64
  }

  func.func @dnrm2(%arg0: memref<?xf64>, %arg1: memref<?xf64>) -> f64  {
    %3 = enzyme.fwddiff @nrm2(%arg0, %arg1) { 
        activity=[#enzyme<activity enzyme_dup>], 
        ret_activity=[#enzyme<activity enzyme_dupnoneed>] 
      } : (memref<?xf64>, memref<?xf64>) -> (f64)
    return %3 : f64
  }

  // CHECK: @fwddiffenrm2(%[[arg0:.+]]: memref<?xf64>, %[[arg1:.+]]: memref<?xf64>) -> f64 {
  // CHECK-DAG:   %[[c0:.+]] = arith.constant 0 : index
  // CHECK-DAG:   %[[c1:.+]] = arith.constant 1 : index
  // CHECK-DAG:   %[[cst:.+]] = arith.constant 0.000000e+00 : f64
  // CHECK-DAG:   %[[dim:.+]] = memref.dim %[[arg0]], %[[c0]] : memref<?xf64>
  // CHECK:   %[[cst_0:.+]] = arith.constant 0.000000e+00 : f64
  // CHECK:   %[[x0:.+]]:2 = scf.parallel (%[[arg2:.+]]) = (%[[c0]]) to (%[[dim]]) step (%[[c1]]) init (%[[cst]], %[[cst_0]]) -> (f64, f64) {
  // CHECK:     %[[x1:.+]] = memref.load %[[arg1]][%[[arg2]]] : memref<?xf64>
  // CHECK:     %[[x2:.+]] = memref.load %[[arg0]][%[[arg2]]] : memref<?xf64>
  // CHECK:     %[[x3:.+]] = arith.mulf %[[x1]], %[[x2]] : f64
  // CHECK:     %[[x4:.+]] = arith.mulf %[[x1]], %[[x2]] : f64
  // CHECK:     %[[x5:.+]] = arith.addf %[[x3]], %[[x4]] : f64
  // CHECK:     %[[x6:.+]] = arith.mulf %[[x2]], %[[x2]] : f64
  // CHECK:     scf.reduce(%[[x6]], %[[x5]] : f64, f64) {
  // CHECK:     ^bb0(%[[arg3:.+]]: f64, %[[arg4:.+]]: f64):
  // CHECK:       %[[x7:.+]] = arith.addf %[[arg3]], %[[arg4]] : f64
  // CHECK:       scf.reduce.return %[[x7]] : f64
  // CHECK:     }, {
  // CHECK:     ^bb0(%[[arg3:.+]]: f64, %[[arg4:.+]]: f64):
  // CHECK:       %[[x7:.+]] = arith.addf %[[arg3]], %[[arg4]] : f64
  // CHECK:       scf.reduce.return %[[x7]] : f64
  // CHECK:     }
  // CHECK:   }
  // CHECK:   return %[[x0]]#1 : f64
  // CHECK: }

}