diffsl 0.9.3

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

// RUN: %eopt --enzyme -canonicalize --remove-unnecessary-enzyme-ops -enzyme-simplify-math -canonicalize %s | FileCheck %s

module {
  func.func @ppow(%x: f64) -> f64 {
    %cst = arith.constant 1.000000e+00 : f64
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %n = arith.constant 10 : index
    %r = scf.for %iv = %c0 to %n step %c1 iter_args(%r_it = %cst) -> f64 {
      %r_next = arith.mulf %r_it, %x : f64
      scf.yield %r_next : f64
    } {enzyme.cache_use_tensor}
    return %r : f64
  }

  func.func @dppow(%x: f64, %dr: f64) -> f64 {
    %r = enzyme.autodiff @ppow(%x, %dr) { activity=[#enzyme<activity enzyme_active>], ret_activity=[#enzyme<activity enzyme_activenoneed>] } : (f64, f64) -> f64
    return %r : f64
  }
}

// CHECK:  func.func private @diffeppow(%[[x:.+]]: f64, %[[dr:.+]]: f64) -> f64 {
// CHECK-NEXT:    %c9 = arith.constant 9 : index
// CHECK-NEXT:    %c10 = arith.constant 10 : index
// CHECK-NEXT:    %c1 = arith.constant 1 : index
// CHECK-NEXT:    %c0 = arith.constant 0 : index
// CHECK-NEXT:    %[[one:.+]] = arith.constant 1.0
// CHECK-NEXT:    %[[zero:.+]] = arith.constant 0.000000e+00 : f64
// CHECK-NEXT:    %[[cache:.+]] = tensor.empty() : tensor<10xf64>
// CHECK-NEXT:    %{{.+}} = scf.for %[[iv:.+]] = %c0 to %c10 step %c1 iter_args(%[[r_it:.+]] = %[[one]], %[[cache_iter:.+]] = %[[cache]]) -> (f64, tensor<10xf64>) {
// CHECK-NEXT:      %[[cache_new:.+]] = tensor.insert %[[r_it]] into %[[cache_iter]][%[[iv]]] : tensor<10xf64>
// CHECK-NEXT:      %[[fwd:.+]] = arith.mulf %[[r_it]], %[[x]] : f64
// CHECK-NEXT:      scf.yield %[[fwd]], %[[cache_new]] : f64, tensor<10xf64>
// CHECK-NEXT:    }
// CHECK-NEXT:    %{{.+}} = scf.for %[[div:.+]] = %c0 to %c10 step %c1 iter_args(%[[dr_it:.+]] = %[[dr]], %[[dx0:.+]] = %[[zero]]) -> (f64, f64) {
// CHECK-NEXT:      %[[rev_idx:.+]] = arith.subi %c9, %[[div]] : index
// CHECK-NEXT:      %[[r_cached:.+]] = tensor.extract %1#1[%[[rev_idx]]] : tensor<10xf64>
// CHECK-NEXT:      %[[dr_next:.+]] = arith.mulf %[[dr_it]], %[[x]] : f64
// CHECK-NEXT:      %[[dx_next:.+]] = arith.mulf %[[dr_it]], %[[r_cached]] : f64
// CHECK-NEXT:      %[[dx1:.+]] = arith.addf %[[dx0]], %[[dx_next]]
// CHECK-NEXT:      scf.yield %[[dr_next]], %[[dx1]] : f64, f64
// CHECK-NEXT:    }
// CHECK-NEXT:    return %2#1 : f64