Introduce patterns provenance-based DUS simplification

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -469,6 +469,27 @@ def UpdateWithoutCornersOp: EnzymeXLA_Op<
   let results = (outs HLO_Tensor:$result);
 }
 
+def PiecewiseSelectOp : EnzymeXLA_Op<"piecewise_select", [Pure, SameOperandsAndResultType]> {
+  let summary = "Piecewise selection operation";
+  let description = [{
+    For coordinates that fit inside any of the boxes defined by the attribute,
+    the result contains the element from the first operand, for other coordinates,
+    it contains the element from the second operand.
+  }];
+
+  let arguments = (ins
+    HLO_Tensor:$whenInBox,
+    HLO_Tensor:$whenOutOfBox,
+    // Boxes is an array where each element is itself an array of integers defining one box.
+    // Each box is defined by 2 * rank values: min (inclusive) and max (non-inclusive) for each dimension, e.g,
+    // [2, 8, 4, 12] defines a box in 2D space with min corner at (2, 4) and max corner at (8, 12).
+    // Alternatively, the box spans [2, 8) along x and [4, 12) along y.
+    ArrayAttr:$boxes
+  );
+
+  let results = (outs HLO_Tensor:$result);
+}
+
 def CommRegionOp : EnzymeXLA_Op<"comm_region", [
     DeclareOpInterfaceMethods<RegionBranchOpInterface>,
     RecursiveMemoryEffects, RecursivelySpeculatable]> {