StaticSelect use analysis

EnzymeAD · Jan 9, 2025 · e881a7e · e881a7e
1 parent ca6aafa
commit e881a7e
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td
@@ -13,7 +13,7 @@ class ConstantFP<string val> : Operation</*primal*/0, /*shadow*/0> {
   string value = val;
 }
 
-class StaticSelect<string condition_> : Operation</*primal*/0, /*shadow*/0, /*custom*/1> {
+class StaticSelect<string condition_> : Operation</*primal*/0, /*shadow*/0, /*custom*/0> {
   string condition = condition_;
 }
 

diff --git a/enzyme/Enzyme/MLIR/Implementations/Common.td b/enzyme/Enzyme/MLIR/Implementations/Common.td
@@ -96,7 +96,7 @@ class GlobalExpr<bit uses_primal, bit uses_shadow, string val> : Operation<uses_
 // The three arguments version accepts a name as a first argument
 // which is then available in the condition as a `Value` under the
 // variable `imVal`.
-class StaticSelect<string condition_> : Operation</*usesPrimal*/0, /*usesShadow*/0, /*usesCustom*/1> {
+class StaticSelect<string condition_> : Operation</*usesPrimal*/0, /*usesShadow*/0, /*usesCustom*/0> {
   string condition = condition_;
 }
 

diff --git a/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp b/enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp
@@ -1166,8 +1166,10 @@ void handleUse(
   bool usesShadow = Def->getValueAsBit("usesShadow");
   bool usesCustom = Def->getValueAsBit("usesCustom");
 
-  // This only concerns instances of StaticSelect for now
-  if (usesCustom) {
+  (void)usesCustom;
+  assert(!usesCustom);
+
+  if (Def->isSubClassOf("StaticSelect")) {
     auto numArgs = resultTree->getNumArgs();
 
     for (int i = numArgs == 3; i < numArgs; ++i) {