Simplify slicebroadcast if no new op is created

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -2870,9 +2870,6 @@ struct SliceBroadcast final
     if (!bcast)
       return failure();
 
-    if (!llvm::hasSingleElement(bcast->getUsers()))
-      return failure();
-
     SmallVector<int64_t> nbcast_idx;
 
     auto preShape = cast<RankedTensorType>(bcast.getOperand().getType());

test/lit_tests/slicebroadcast.mlir

@@ -13,6 +13,12 @@ module {
     %2389 = stablehlo.slice %2388 [0:1, 0:8, 0:3, 0:1024, 1024:2048] : (tensor<1x8x3x1024x2048xf32>) -> tensor<1x8x3x1024x1024xf32>
     return %2389 : tensor<1x8x3x1024x1024xf32>
   }
+  
+  func.func @main3(%20: tensor<f64>) -> (tensor<3056x6128xf64>, tensor<3055x6128xf64>) {
+    %28 = stablehlo.broadcast_in_dim %20, dims = [] : (tensor<f64>) -> tensor<3056x6128xf64>
+    %31 = stablehlo.slice %28 [1:3056, 0:6128] : (tensor<3056x6128xf64>) -> tensor<3055x6128xf64> 
+    return %28, %31 : tensor<3056x6128xf64>, tensor<3055x6128xf64>
+  }
 }
 
 // CHECK:  func.func @main(%arg0: tensor<2x3x50xf32>) -> tensor<4x1x25x15x2x3xf32> {
@@ -25,3 +31,10 @@ module {
 // CHECK-NEXT:    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1, 2, 3, 4] : (tensor<1x8x3x1024x1xf32>) -> tensor<1x8x3x1024x1024xf32>
 // CHECK-NEXT:    return %0 : tensor<1x8x3x1024x1024xf32>
 // CHECK-NEXT:  }
+
+// CHECK:  func.func @main3(%arg0: tensor<f64>) -> (tensor<3056x6128xf64>, tensor<3055x6128xf64>) {
+// CHECK-NEXT:    %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f64>) -> tensor<3056x6128xf64>
+// CHECK-NEXT:    %1 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f64>) -> tensor<3055x6128xf64>
+// CHECK-NEXT:    return %0, %1 : tensor<3056x6128xf64>, tensor<3055x6128xf64>
+// CHECK-NEXT:  }
+