drop the extend if it becomes a no-op.

Author: jumerckx

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -18771,10 +18771,6 @@ struct ExtendSlice final
     if (!extendOp)
       return rewriter.notifyMatchFailure(op, "Operand is not an ExtendOp");
 
-    if (extendOp.getResult().getNumUses() > 1)
-      return rewriter.notifyMatchFailure(
-          op, "ExtendOp result is used multiple times");
-
     // This transformation is simplified if strides are 1.
     if (llvm::any_of(op.getStrides(), [](int64_t s) { return s != 1; }))
       return rewriter.notifyMatchFailure(op, "Requires strides of 1");
@@ -18811,6 +18807,16 @@ struct ExtendSlice final
     int64_t new_rhs =
         std::max((int64_t)0, limit_d - std::max(start_d, lhs + size_d));
 
+    if (new_lhs == 0 && new_rhs == 0) {
+      auto newSlice = rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
+        op, op.getType(), operand, new_starts, new_limits, new_strides);
+      return success();
+    }
+
+    if (extendOp.getResult().getNumUses() > 1) {
+      return rewriter.notifyMatchFailure(op, "ExtendOp result is used multiple times");
+    }
+
     // Create the new slice on the original tensor.
     auto newSlice = rewriter.create<stablehlo::SliceOp>(
         op.getLoc(), operand, new_starts, new_limits, new_strides);

test/lit_tests/extendslice.mlir

@@ -24,3 +24,16 @@ func.func @f_multiple_uses(%a: tensor<4x1520x3056xf64>) -> (tensor<3x1520x3056xf
   %d = stablehlo.slice %b [0:4, 0:1520, 0:3056] : (tensor<5x1520x3056xf64>) -> tensor<4x1520x3056xf64>
   return %c, %d : tensor<3x1520x3056xf64>, tensor<4x1520x3056xf64>
 }
+
+// CHECK:      func.func @f_multiple_uses_superfluous_extend(%arg0: tensor<4x1520x3056xf64>) -> (tensor<3x1520x3056xf64>, tensor<4x1520x3056xf64>) {
+// CHECK-NEXT:   %0 = stablehlo.slice %arg0 [0:3, 0:1520, 0:3056] : (tensor<4x1520x3056xf64>) -> tensor<3x1520x3056xf64>
+// CHECK-NEXT:   %1 = "enzymexla.extend"(%0) <{dimension = 0 : i64, lhs = 1 : i64, rhs = 0 : i64}> : (tensor<3x1520x3056xf64>) -> tensor<4x1520x3056xf64>
+// CHECK-NEXT:   %2 = stablehlo.slice %arg0 [0:3, 0:1520, 0:3056] : (tensor<4x1520x3056xf64>) -> tensor<3x1520x3056xf64>
+// CHECK-NEXT:   return %2, %1 : tensor<3x1520x3056xf64>, tensor<4x1520x3056xf64>
+// CHECK-NEXT: }
+func.func @f_multiple_uses_superfluous_extend(%a: tensor<4x1520x3056xf64>) -> (tensor<3x1520x3056xf64>, tensor<4x1520x3056xf64>) {
+  %b = "enzymexla.extend"(%a) <{dimension = 0 : i64, lhs = 1 : i64, rhs = 0 : i64}> : (tensor<4x1520x3056xf64>) -> tensor<5x1520x3056xf64>
+  %c = stablehlo.slice %b [0:4, 0:1520, 0:3056] : (tensor<5x1520x3056xf64>) -> tensor<4x1520x3056xf64>
+  %d = stablehlo.slice %b [1:4, 0:1520, 0:3056] : (tensor<5x1520x3056xf64>) -> tensor<3x1520x3056xf64>
+  return %d, %c : tensor<3x1520x3056xf64>, tensor<4x1520x3056xf64>
+}