MLIR: Fix a couple of issues with reverse mode applied to scf.if + scf.parallel (#2723)

* Implement MemRefAutoDiffTypeInterface::createNullValue

Needed for reverse mode of scf.if.  I implemented this by creating a
memref of the given type with 0 as the size for all dynamic dimensions.

* Fix ParallelOpEnzymeOpsRemover::getCanonicalLoopIVs

The step was being subtracted from the nonzero lower bound instead of
the lower bound itself

* Add reverse mode test if scf.if + scf.parallel

* Run clang-format

Author: etphipp

enzyme/Enzyme/MLIR/Implementations/MemRefAutoDiffOpInterfaceImpl.cpp

@@ -242,7 +242,15 @@ class MemRefAutoDiffTypeInterface
   }
   mlir::Value createNullValue(mlir::Type self, OpBuilder &builder,
                               Location loc) const {
-    llvm_unreachable("Cannot create null of memref (todo polygeist null)");
+    // Create a memref of the given type with the required number of
+    // dynamic dimensions, all set to 0
+    MemRefType MT = cast<MemRefType>(self);
+    unsigned numDynamicDims = MT.getNumDynamicDims();
+    SmallVector<mlir::Value> dynamicSizes(numDynamicDims);
+    for (unsigned i = 0; i < numDynamicDims; ++i) {
+      dynamicSizes[i] = builder.create<mlir::arith::ConstantIndexOp>(loc, 0);
+    }
+    return mlir::memref::AllocOp::create(builder, loc, MT, dynamicSizes);
   }
 
   Value createAddOp(Type self, OpBuilder &builder, Location loc, Value a,

enzyme/Enzyme/MLIR/Implementations/SCFAutoDiffOpInterfaceImpl.cpp

@@ -675,7 +675,7 @@ struct ParallelOpEnzymeOpsRemover
                          parOp.getStep())) {
       Value val = iv;
       if (!matchPattern(lb, m_Zero())) {
-        val = arith::SubIOp::create(builder, parOp.getLoc(), val, step);
+        val = arith::SubIOp::create(builder, parOp.getLoc(), val, lb);
       }
 
       if (!matchPattern(step, m_One())) {

enzyme/test/MLIR/ReverseMode/scf_parallel_if.mlir

@@ -0,0 +1,60 @@
+// RUN: %eopt %s --enzyme --canonicalize --remove-unnecessary-enzyme-ops --canonicalize --enzyme-simplify-math --cse | FileCheck %s
+
+module {
+  func.func @scale(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: memref<?xf64>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %dim = memref.dim %arg0, %c0 : memref<?xf64>
+    %0 = arith.cmpi sgt, %dim, %c0 : index
+    scf.if %0 {
+      scf.parallel (%arg3) = (%c0) to (%dim) step (%c1) {
+        %1 = memref.load %arg0[%arg3] : memref<?xf64>
+        %2 = memref.load %arg1[%arg3] : memref<?xf64>
+        %3 = arith.mulf %1, %2 : f64
+        memref.store %3, %arg2[%arg3] : memref<?xf64>
+      }
+    }
+    return
+  }
+
+  func.func @dscale(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: memref<?xf64>, %arg3: memref<?xf64>, %arg4: memref<?xf64>) {
+    enzyme.autodiff @scale(%arg0, %arg3, %arg1, %arg2, %arg4) {
+        activity=[#enzyme<activity enzyme_dup>, 
+                  #enzyme<activity enzyme_const>, 
+                  #enzyme<activity enzyme_dup>],
+        ret_activity=[]
+      } : (memref<?xf64>, memref<?xf64>, memref<?xf64>, memref<?xf64>, memref<?xf64>) -> ()
+
+    return
+  }
+
+  // CHECK: @diffescale(%[[arg0:.+]]: memref<?xf64>, %[[arg1:.+]]: memref<?xf64>, %[[arg2:.+]]: memref<?xf64>, %[[arg3:.+]]: memref<?xf64>, %[[arg4:.+]]: memref<?xf64>) {
+  // CHECK:   %[[c1:.+]] = arith.constant 1 : index
+  // CHECK:   %[[c0:.+]] = arith.constant 0 : index
+  // CHECK:   %[[cst:.+]] = arith.constant 0.000000e+00 : f64
+  // CHECK:   %[[dim:.+]] = memref.dim %[[arg0]], %[[c0]] : memref<?xf64>
+  // CHECK:   %[[x0:.+]] = arith.cmpi sgt, %[[dim]], %[[c0]] : index
+  // CHECK:   scf.if %[[x0]] {
+  // CHECK:     %[[alloc:.+]] = memref.alloc(%[[dim]]) : memref<?xf64>
+  // CHECK:     scf.parallel (%[[arg5:.+]]) = (%[[c0]]) to (%[[dim]]) step (%[[c1]]) {
+  // CHECK:       %[[x1:.+]] = memref.load %arg0[%[[arg5]]] : memref<?xf64>
+  // CHECK:       %[[x2:.+]] = memref.load %arg2[%[[arg5]]] : memref<?xf64>
+  // CHECK:       memref.store %[[x2]], %[[alloc]][%[[arg5]]] : memref<?xf64>
+  // CHECK:       %[[x3:.+]] = arith.mulf %[[x1]], %[[x2]] : f64
+  // CHECK:       memref.store %[[x3]], %arg3[%[[arg5]]] : memref<?xf64>
+  // CHECK:       scf.reduce 
+  // CHECK:     }
+  // CHECK:     scf.parallel (%[[arg5:.+]]) = (%[[c0]]) to (%[[dim]]) step (%[[c1]]) {
+  // CHECK:       %[[x1:.+]] = memref.load %[[alloc]][%[[arg5]]] : memref<?xf64>
+  // CHECK:       %[[x2:.+]] = memref.load %arg4[%[[arg5]]] : memref<?xf64>
+  // CHECK:       memref.store %[[cst]], %arg4[%[[arg5]]] : memref<?xf64>
+  // CHECK:       %[[x3:.+]] = arith.mulf %[[x2]], %[[x1]] : f64
+  // CHECK:       %[[x4:.+]] = memref.atomic_rmw addf %[[x3]], %arg1[%[[arg5]]] : (f64, memref<?xf64>) -> f64
+  // CHECK:       scf.reduce 
+  // CHECK:     }
+  // CHECK:     memref.dealloc %[[alloc]] : memref<?xf64>
+  // CHECK:   }
+  // CHECK:   return
+  // CHECK: }
+
+}