Mark recursively speculatable ops as such (openxla#2175)

This includes the regioned ops case, if and while.
ghpvnist · Apr 8, 2024 · 5e03518 · 5e03518
1 parent e2ad824
commit 5e03518
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td
@@ -1197,6 +1197,7 @@ def StableHLO_AfterAllOp : StableHLO_Op<"after_all", [Pure,
 // conditional use of kConditional HLO.
 def StableHLO_IfOp: StableHLO_Op<"if", [
     RecursiveMemoryEffects,
+    RecursivelySpeculatable,
     SingleBlockImplicitTerminator<"ReturnOp">,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "If operation";
@@ -1230,6 +1231,7 @@ def StableHLO_IfOp: StableHLO_Op<"if", [
 // conditional use of kConditional HLO.
 def StableHLO_CaseOp: StableHLO_Op<"case", [
       RecursiveMemoryEffects,
+      RecursivelySpeculatable,
       SingleBlockImplicitTerminator<"ReturnOp">,
       DeclareOpInterfaceMethods<InferTypeOpInterface> /*case_c4*/
     ]> {
@@ -1263,6 +1265,7 @@ def StableHLO_CaseOp: StableHLO_Op<"case", [
 
 def StableHLO_WhileOp: StableHLO_Op<"while", [
       RecursiveMemoryEffects,
+      RecursivelySpeculatable,
       SingleBlockImplicitTerminator<"ReturnOp">,
       DeclareOpInterfaceMethods<InferTypeOpInterface> /*while_c3*/,
       OpAsmOpInterface