verify rule validity

Author: Pangoraw

enzyme/Enzyme/MLIR/Dialect/EnzymeOps.td

@@ -206,6 +206,11 @@ def CustomReverseRuleOp : Enzyme_Op<"custom_reverse_rule", [IsolatedFromAbove, S
   let assemblyFormat = [{
     $sym_name $body attr-dict-with-keyword
   }];
+
+  let extraClassDeclaration = [{
+    llvm::LogicalResult activityMatch(llvm::ArrayRef<enzyme::Activity> ArgActivity,
+                                      llvm::ArrayRef<enzyme::Activity> RetActivity);
+  }];
 }
 
 def CustomReverseRuleAugmentedPrimalOp : Enzyme_Op<"custom_reverse_rule.augmented_primal", [

enzyme/Enzyme/MLIR/Dialect/Ops.cpp

@@ -913,6 +913,39 @@ static void printAugmentedFn(OpAsmPrinter &p, FunctionType fnType,
                 /*printBlockTerminators*/ true);
 }
 
+//===----------------------------------------------------------------------===//
+// CustomReverseRuleOp
+//===----------------------------------------------------------------------===//
+
+llvm::LogicalResult CustomReverseRuleOp::activityMatch(
+    llvm::ArrayRef<enzyme::Activity> argActivity,
+    llvm::ArrayRef<enzyme::Activity> retActivity) {
+  auto selfArgActivity = getActivity();
+  auto selfRetActivity = getRetActivity();
+
+  if (selfArgActivity.size() != argActivity.size() ||
+      selfRetActivity.size() != retActivity.size())
+    return failure();
+
+  for (auto [attr, act] : llvm::zip_equal(selfArgActivity, argActivity)) {
+    auto iattr = cast<ActivityAttr>(attr);
+    auto val = iattr.getValue();
+
+    if (val == Activity::enzyme_const && act != Activity::enzyme_const)
+      return failure();
+  }
+
+  for (auto [attr, act] : llvm::zip_equal(selfRetActivity, retActivity)) {
+    auto iattr = cast<ActivityAttr>(attr);
+    auto val = iattr.getValue();
+
+    if (val == Activity::enzyme_const && act != Activity::enzyme_const)
+      return failure();
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CustomReverseRuleAugmentedPrimalOp
 //===----------------------------------------------------------------------===//

enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp

@@ -692,6 +692,33 @@ static LogicalResult callReverseHandlerSplit(Operation *orig,
     // point it at the rule the callee named.
     cr = cast<CustomReverseRuleOp>(crOp);
     primalCall = cast<CallAugmentedPrimalOp>(gutils->getNewFromOriginal(orig));
+
+    auto diffeTypeToActivity = [](DIFFE_TYPE act) {
+      switch (act) {
+      case DIFFE_TYPE::CONSTANT:
+        return Activity::enzyme_const;
+      case DIFFE_TYPE::OUT_DIFF:
+        return Activity::enzyme_active;
+      case DIFFE_TYPE::DUP_ARG:
+        return Activity::enzyme_dup;
+      case DIFFE_TYPE::DUP_NONEED:
+        return Activity::enzyme_dupnoneed;
+      default:
+        llvm_unreachable("cannot handle act");
+      }
+    };
+
+    SmallVector<enzyme::Activity> ArgActivityAct =
+        llvm::map_to_vector(ArgActivity, diffeTypeToActivity);
+    SmallVector<enzyme::Activity> RetActivityAct =
+        llvm::map_to_vector(RetActivity, diffeTypeToActivity);
+
+    if (failed(cr.activityMatch(ArgActivityAct, RetActivityAct)))
+      return orig->emitError()
+             << "could not find a rule with the right activity (rule activity="
+             << cr.getActivity() << ", ret_activity=" << cr.getRetActivity()
+             << ")";
+
   } else {
     std::vector<bool> overwritten_args(narg, true);
     std::vector<bool> returnShadow(nret, false);

enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp

@@ -331,13 +331,41 @@ FlatSymbolRefAttr MEnzymeLogic::CreateSplitModeDiff(
 
   SymbolTable symbolTable(SymbolTable::getNearestSymbolTable(fn));
 
+  SmallVector<mlir::Attribute> argAttrs;
+  if (auto prevArgAttrs = fn.getAllArgAttrs())
+    argAttrs.assign(prevArgAttrs.begin(), prevArgAttrs.end());
+
+  SmallVector<Attribute> argActivityAttrs;
+  for (auto [i, act] : llvm::enumerate(constants)) {
+    argActivityAttrs.push_back(activityFromDiffeType(fn.getContext(), act));
+
+    if (!argAttrs.empty() && act == DIFFE_TYPE::DUP_ARG)
+      argAttrs.insert(argAttrs.begin() + i + 1 -
+                          (argAttrs.size() - fn.getNumArguments()),
+                      nullptr);
+  }
+
+  SmallVector<Attribute> retActivityAttrs;
+  for (auto act : retType)
+    retActivityAttrs.push_back(activityFromDiffeType(fn.getContext(), act));
+
   if (auto existingCustomRule =
           fn->getAttrOfType<FlatSymbolRefAttr>("enzyme.custom_rule")) {
     auto CR = symbolTable.lookup<enzyme::CustomReverseRuleOp>(
         existingCustomRule.getValue());
 
     if (CR) {
-      return existingCustomRule;
+      auto getAttrActivity = [](auto attr) {
+        return cast<ActivityAttr>(attr).getValue();
+      };
+
+      SmallVector<Activity> ArgActivity =
+          llvm::map_to_vector(argActivityAttrs, getAttrActivity);
+      SmallVector<Activity> RetActivity =
+          llvm::map_to_vector(retActivityAttrs, getAttrActivity);
+
+      if (!failed(CR.activityMatch(ArgActivity, RetActivity)))
+        return existingCustomRule;
     }
   }
 
@@ -360,24 +388,6 @@ FlatSymbolRefAttr MEnzymeLogic::CreateSplitModeDiff(
 
   auto name = fn.getName();
 
-  SmallVector<mlir::Attribute> argAttrs;
-  if (auto prevArgAttrs = fn.getAllArgAttrs())
-    argAttrs.assign(prevArgAttrs.begin(), prevArgAttrs.end());
-
-  SmallVector<Attribute> argActivityAttrs;
-  for (auto [i, act] : llvm::enumerate(constants)) {
-    argActivityAttrs.push_back(activityFromDiffeType(fn.getContext(), act));
-
-    if (!argAttrs.empty() && act == DIFFE_TYPE::DUP_ARG)
-      argAttrs.insert(argAttrs.begin() + i + 1 -
-                          (argAttrs.size() - fn.getNumArguments()),
-                      nullptr);
-  }
-
-  SmallVector<Attribute> retActivityAttrs;
-  for (auto act : retType)
-    retActivityAttrs.push_back(activityFromDiffeType(fn.getContext(), act));
-
   auto argActivityAttr = ArrayAttr::get(fn.getContext(), argActivityAttrs);
   auto retActivityAttr = ArrayAttr::get(fn.getContext(), retActivityAttrs);
 

enzyme/test/MLIR/ReverseMode/need_rule_set.mlir

@@ -0,0 +1,45 @@
+// RUN: %eopt %s --enzyme-wrap="infn=main outfn= argTys=enzyme_active,enzyme_active retTys=enzyme_active mode=ReverseModeCombined" --verify-diagnostics
+
+module {
+  enzyme.custom_reverse_rule @reverse_f {
+    %cache_exp = "enzyme.init"() : () -> !enzyme.Cache<f32>
+    %cache_arg1 = "enzyme.init"() : () -> !enzyme.Cache<f32>
+
+    enzyme.custom_reverse_rule.augmented_primal (%arg0: f32, %arg1: f32) -> f32 {
+      "enzyme.push"(%cache_arg1, %arg1) : (!enzyme.Cache<f32>, f32) -> ()
+      %0 = math.exp %arg0 : f32
+      "enzyme.push"(%cache_exp, %0) : (!enzyme.Cache<f32>, f32) -> ()
+      %1 = arith.mulf %arg1, %arg0 : f32
+      enzyme.yield %1 : f32
+    }
+
+    enzyme.custom_reverse_rule.augmented_primal (%dres: f32) -> f32 {
+      %exp = "enzyme.pop"(%cache_exp) : (!enzyme.Cache<f32>) -> f32
+      %arg1 = "enzyme.pop"(%cache_exp) : (!enzyme.Cache<f32>) -> f32
+      %d1 = arith.mulf %dres, %arg1 : f32
+      %darg0 = arith.mulf %exp, %d1 : f32
+      enzyme.yield %darg0 : f32
+    }
+
+    enzyme.yield
+  } attributes {
+    activity=[#enzyme<activity enzyme_active>,
+              #enzyme<activity enzyme_const>],
+    ret_activity=[#enzyme<activity enzyme_active>],
+    function_type = (f32, f32) -> f32
+  }
+
+  func.func @f(%arg0:f32, %arg1: f32) -> f32 attributes {enzyme.custom_rule = @reverse_f} {
+    %0 = math.exp %arg0 : f32
+    %1 = arith.mulf %arg1, %arg0 : f32
+    return %1 : f32
+  }
+
+  func.func @main(%arg0: f32, %arg1: f32) -> f32 {
+
+    // expected-error @below {{could not find a rule with the right activity (rule activity=[#enzyme<activity enzyme_active>, #enzyme<activity enzyme_const>], ret_activity=[#enzyme<activity enzyme_active>])}}
+    %0 = func.call @f( %arg0, %arg1 ) : (f32, f32) -> f32
+
+    return %0 : f32
+  }
+}