feat: expand coverage of noop_reverse (#1568)

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -13090,6 +13090,12 @@ struct NoopReverse final
 
   LogicalResult matchAndRewriteImpl(stablehlo::ReverseOp op,
                                     PatternRewriter &rewriter) const {
+    SplatElementsAttr splat;
+    if (matchPattern(op.getOperand(), m_Constant(&splat))) {
+      rewriter.replaceAllUsesWith(op, op.getOperand());
+      return success();
+    }
+
     SmallVector<int64_t> newDimensions;
     auto dimensions = op.getDimensions();
     auto shape = op.getResult().getType().getShape();
@@ -13100,6 +13106,11 @@ struct NoopReverse final
         newDimensions.push_back(dim);
     }
 
+    if (auto bcast =
+            op.getOperand().getDefiningOp<stablehlo::BroadcastInDimOp>()) {
+      peelBroadcastedDimensions(bcast, newDimensions);
+    }
+
     if (newDimensions.empty()) {
       rewriter.replaceOp(op, op.getOperand());
       return success();
@@ -13112,6 +13123,27 @@ struct NoopReverse final
                                                       newDimensions);
     return success();
   }
+
+private:
+  void peelBroadcastedDimensions(stablehlo::BroadcastInDimOp op,
+                                 SmallVectorImpl<int64_t> &dims) const {
+    DenseMap<int64_t, int64_t> dimMap;
+    for (auto [i, dim] : llvm::enumerate(op.getBroadcastDimensions())) {
+      dimMap[dim] = i;
+    }
+
+    auto opShape = cast<RankedTensorType>(op.getOperand().getType()).getShape();
+
+    auto newEnd = llvm::remove_if(dims, [&](int64_t dim) {
+      auto it = dimMap.find(dim);
+      if (it != dimMap.end()) {
+        return opShape[it->second] == 1; // if 1 then trivially expanded
+      }
+      return true; // not in broadcast dims so it was expanded
+    });
+    dims.erase(newEnd, dims.end());
+    return;
+  }
 };
 
 /// Converts gather ops to slice ops in case we have a single set of constant

test/lit_tests/reverse_simplify.mlir

@@ -0,0 +1,45 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-generate-td="patterns=noop_reverse" --transform-interpreter --enzyme-hlo-remove-transform %s | FileCheck %s
+
+module {
+  func.func @main() -> tensor<8x4x3xf32> {
+    %cst = stablehlo.constant dense<0.000000e+00> : tensor<8x4x3xf32>
+    %1 = stablehlo.reverse %cst, dims = [2, 1] : tensor<8x4x3xf32>
+    return %1 : tensor<8x4x3xf32>
+    // CHECK: %cst = stablehlo.constant dense<0.000000e+00> : tensor<8x4x3xf32>
+    // CHECK-NEXT: return %cst : tensor<8x4x3xf32>
+  }
+}
+
+module {
+  func.func @main(%arg0: tensor<8x1xf32>) -> tensor<8x4x3x1xf32> {
+    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 3] : (tensor<8x1xf32>) -> tensor<8x4x3x1xf32>
+    %1 = stablehlo.reverse %0, dims = [3, 2, 0] : tensor<8x4x3x1xf32>
+    return %1 : tensor<8x4x3x1xf32>
+
+    // CHECK: %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 3] : (tensor<8x1xf32>) -> tensor<8x4x3x1xf32>
+    // CHECK-NEXT: %1 = stablehlo.reverse %0, dims = [0] : tensor<8x4x3x1xf32>
+    // CHECK-NEXT: return %1 : tensor<8x4x3x1xf32>
+  }
+}
+
+module {
+  func.func @main(%arg0: tensor<8x1xf32>) -> tensor<8x4x3x1xf32> {
+    %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 3] : (tensor<8x1xf32>) -> tensor<8x4x3x1xf32>
+    %1 = stablehlo.reverse %0, dims = [2, 0] : tensor<8x4x3x1xf32>
+    return %1 : tensor<8x4x3x1xf32>
+
+    // CHECK: %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 3] : (tensor<8x1xf32>) -> tensor<8x4x3x1xf32>
+    // CHECK-NEXT: %1 = stablehlo.reverse %0, dims = [0] : tensor<8x4x3x1xf32>
+    // CHECK-NEXT: return %1 : tensor<8x4x3x1xf32>
+  }
+}
+
+module {
+  func.func @main(%arg0: tensor<1x8xf32>) -> tensor<1x8xf32> {
+    %0 = stablehlo.reverse %arg0, dims = [0, 1] : tensor<1x8xf32>
+    return %0 : tensor<1x8xf32>
+
+    // CHECK: %0 = stablehlo.reverse %arg0, dims = [1] : tensor<1x8xf32>
+    // CHECK-NEXT: return %0 : tensor<1x8xf32>
+  }
+}