Fix update (#1849)

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -2361,7 +2361,9 @@ struct DUSDUSConcat final
                idxs[0] < idxs[1] + tys[1].getShape()[diffidx] &&
                idxs[0] + tys[0].getShape()[diffidx] >
                    idxs[1] + tys[1].getShape()[diffidx] &&
-               allStatic) {
+               allStatic &&
+               idxs[0] + tys[0].getShape()[diffidx] <=
+                   dus.getOperand().getType().getShape()[diffidx]) {
       // the new update overlaps, following the old update
 
       // Case 5:
@@ -2386,7 +2388,9 @@ struct DUSDUSConcat final
                idxs[1] < idxs[0] + tys[0].getShape()[diffidx] &&
                idxs[0] + tys[0].getShape()[diffidx] <
                    idxs[1] + tys[1].getShape()[diffidx] &&
-               allStatic) {
+               allStatic &&
+               idxs[1] + tys[1].getShape()[diffidx] <=
+                   dus.getOperand().getType().getShape()[diffidx]) {
       // the new update overlaps, following the old update
 
       // Case 5:

test/lit_tests/dusdus2.mlir

@@ -0,0 +1,20 @@
+// RUN: enzymexlamlir-opt %s --enzyme-hlo-opt --split-input-file | FileCheck %s
+
+module {
+  func.func @fuse(%op: tensor<4x1x8x8xcomplex<f32>>, %19: tensor<4x1x4x8xcomplex<f32>>, %22: tensor<4x1x4x8xcomplex<f32>>) -> tensor<4x1x8x8xcomplex<f32>> {
+    %c_7 = stablehlo.constant dense<4> : tensor<i32>
+    %c_6 = stablehlo.constant dense<0> : tensor<i32>
+    %23 = stablehlo.dynamic_update_slice %op, %19, %c_6, %c_6, %c_6, %c_6 : (tensor<4x1x8x8xcomplex<f32>>, tensor<4x1x4x8xcomplex<f32>>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4x1x8x8xcomplex<f32>>
+    %26 = stablehlo.dynamic_update_slice %23, %22, %c_6, %c_6, %c_6, %c_7 : (tensor<4x1x8x8xcomplex<f32>>, tensor<4x1x4x8xcomplex<f32>>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4x1x8x8xcomplex<f32>>
+    func.return %26: tensor<4x1x8x8xcomplex<f32>>
+  }
+}
+
+// CHECK:  func.func @fuse(%arg0: tensor<4x1x8x8xcomplex<f32>>, %arg1: tensor<4x1x4x8xcomplex<f32>>, %arg2: tensor<4x1x4x8xcomplex<f32>>) -> tensor<4x1x8x8xcomplex<f32>> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<4> : tensor<i32>
+// CHECK-NEXT:    %c_0 = stablehlo.constant dense<0> : tensor<i32>
+// CHECK-NEXT:    %0 = stablehlo.slice %arg0 [0:4, 0:1, 4:8, 0:8] : (tensor<4x1x8x8xcomplex<f32>>) -> tensor<4x1x4x8xcomplex<f32>>
+// CHECK-NEXT:    %1 = stablehlo.dynamic_update_slice %arg1, %arg2, %c_0, %c_0, %c_0, %c : (tensor<4x1x4x8xcomplex<f32>>, tensor<4x1x4x8xcomplex<f32>>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4x1x4x8xcomplex<f32>>
+// CHECK-NEXT:    %2 = stablehlo.concatenate %1, %0, dim = 2 : (tensor<4x1x4x8xcomplex<f32>>, tensor<4x1x4x8xcomplex<f32>>) -> tensor<4x1x8x8xcomplex<f32>>
+// CHECK-NEXT:    return %2 : tensor<4x1x8x8xcomplex<f32>>
+// CHECK-NEXT:  }