feat: partial support for broadcast elementwise trait (#1602)

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3086,7 +3086,7 @@ struct SliceElementwise final
     auto elem = op.getOperand().getDefiningOp();
     if (!elem)
       return failure();
-    if (!elem->hasTrait<mlir::OpTrait::Elementwise>())
+    if (!stablehlo::hasTraitElementwise(elem))
       return failure();
     if (llvm::hasSingleElement(elem->getUsers())) {
       SmallVector<Value> ops;
@@ -3772,7 +3772,7 @@ struct FullReduceReshapeOrTranspose final
           reshapeOrTransposes.push_back(rs);
           continue;
         }
-        if (!curOp->hasTrait<mlir::OpTrait::Elementwise>())
+        if (!hasTraitElementwise(curOp))
           return failure();
         if (!isMemoryEffectFree(curOp))
           return failure();
@@ -10469,7 +10469,7 @@ struct SliceReshapeElementwise final
     auto elem = reshape.getOperand().getDefiningOp();
     if (!elem)
       return failure();
-    if (!elem->hasTrait<mlir::OpTrait::Elementwise>())
+    if (!stablehlo::hasTraitElementwise(elem))
       return failure();
     if (!llvm::hasSingleElement(elem->getUsers()))
       return failure();
@@ -10515,7 +10515,7 @@ struct TransposeElementwise final
     if (!elem)
       return failure();
 
-    if (!elem->hasTrait<mlir::OpTrait::Elementwise>())
+    if (!stablehlo::hasTraitElementwise(elem))
       return failure();
 
     bool singleUser = llvm::hasSingleElement(elem->getUsers());
@@ -10913,7 +10913,7 @@ struct ReshapeElementwise final
     if (onlySingleUser && !llvm::hasSingleElement(elem->getUsers()))
       return failure();
 
-    if (!elem->hasTrait<mlir::OpTrait::Elementwise>())
+    if (!stablehlo::hasTraitElementwise(elem))
       return failure();
 
     SmallVector<Value> ops;
@@ -18958,8 +18958,8 @@ LogicalResult commUnaryOpElementwise(bool onlySingleUser, EnzymeOp op,
 
   bool anyModified = false;
   for (auto elem : llvm::make_early_inc_range(op->getUsers())) {
-    if (!elem->template hasTrait<mlir::OpTrait::Elementwise>() ||
-        elem->getNumResults() != 1 || elem->getNumOperands() != 1)
+    if (!hasTraitElementwise(elem) || elem->getNumResults() != 1 ||
+        elem->getNumOperands() != 1)
       continue;
 
     auto newOp = rewriter.create(
@@ -20270,7 +20270,7 @@ struct ConcatElementwise final
       if (isa<stablehlo::ConvertOp>(vdefOp)) // Conflicts with ConvertConcat
         return failure();
 
-      if (vdefOp->hasTrait<mlir::OpTrait::Elementwise>()) {
+      if (hasTraitElementwise(vdefOp)) {
         if (concatOpOperands.size() != 0) {
           if (!OperationEquivalence::isEquivalentTo(
                   concatOpOperands[0], vdefOp,
@@ -20671,7 +20671,7 @@ struct ConcatReshapeElementwise final
       if (!vdefOp)
         return failure();
 
-      if (vdefOp->hasTrait<mlir::OpTrait::Elementwise>()) {
+      if (hasTraitElementwise(vdefOp)) {
         if (concatOpOperands.size() != 0) {
           if (!OperationEquivalence::isEquivalentTo(
                   concatOpOperands[0], vdefOp,
@@ -21259,7 +21259,7 @@ struct GatherElementwise
                                     PatternRewriter &rewriter) const {
     auto gatherInput = op.getOperand();
     auto defOp = gatherInput.getDefiningOp();
-    if (!defOp || !defOp->hasTrait<mlir::OpTrait::Elementwise>())
+    if (!defOp || !hasTraitElementwise(defOp))
       return rewriter.notifyMatchFailure(op,
                                          "GatherOp with non-elementwise input");
 

src/enzyme_ad/jax/Utils.cpp

@@ -1380,6 +1380,38 @@ Value reshapeAxisOutOf(OpBuilder &builder, Value input,
       builder.getDenseI64ArrayAttr(permutation));
 }
 
+bool hasTraitElementwise(Operation *op) {
+  if (op->hasTrait<OpTrait::Elementwise>())
+    return true;
+
+  if (op->hasTrait<hlo::OpTrait::BroadcastingElementwise>()) {
+    // Check sizes (shapes) match across operands, not the exact types.
+    auto refShapedTy = dyn_cast<ShapedType>(op->getOperand(0).getType());
+    if (!refShapedTy)
+      return false;
+
+    for (auto operand : op->getOperands()) {
+      auto curShapedTy = dyn_cast<ShapedType>(operand.getType());
+      if (!curShapedTy)
+        return false;
+
+      if (curShapedTy.getRank() != refShapedTy.getRank())
+        return false;
+
+      for (int64_t i = 0; i < curShapedTy.getRank(); ++i) {
+        int64_t a = curShapedTy.getDimSize(i);
+        int64_t b = refShapedTy.getDimSize(i);
+        // If both are static and different, sizes don't match.
+        if (a != ShapedType::kDynamic && b != ShapedType::kDynamic && a != b)
+          return false;
+      }
+    }
+    return true;
+  }
+
+  return false;
+}
+
 } // namespace stablehlo
 
 } // namespace mlir

src/enzyme_ad/jax/Utils.h

@@ -798,6 +798,11 @@ mlir::Value reshapeAxisInto(OpBuilder &builder, Value input,
 mlir::Value reshapeAxisOutOf(OpBuilder &builder, Value input,
                              ArrayRef<int64_t> &batchSizes, int64_t dim);
 
+// matches for hasTrait<OpTrait::Elementwise>. Additionally matches for
+// hasTrait<OpTrait::HLOBroadcastingElementwise> if all of the operands are
+// of the same shape.
+bool hasTraitElementwise(Operation *op);
+
 } // namespace stablehlo
 
 } // namespace mlir

test/lit_tests/transpose_clamp.mlir

@@ -0,0 +1,20 @@
+// RUN: enzymexlamlir-opt  --enzyme-hlo-generate-td="patterns=transpose_elementwise(0);transpose_simplify;transpose_transpose" --transform-interpreter --enzyme-hlo-remove-transform %s | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<4x3xf32>) -> tensor<4x3xf32> {
+    %cst = stablehlo.constant dense<6.000000e+00> : tensor<3x4xf32>
+    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<3x4xf32>
+    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x3xf32>) -> tensor<3x4xf32>
+    %1 = stablehlo.clamp %cst_0, %0, %cst : tensor<3x4xf32>
+    %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<3x4xf32>) -> tensor<4x3xf32>
+    return %2 : tensor<4x3xf32>
+  }
+}
+
+// CHECK: func.func @main(%arg0: tensor<4x3xf32>) -> tensor<4x3xf32> {
+// CHECK-NEXT:   %cst = stablehlo.constant dense<6.000000e+00> : tensor<4x3xf32>
+// CHECK-NEXT:   %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<4x3xf32>
+// CHECK-NEXT:   %0 = stablehlo.transpose %arg0, dims = [0, 1] : (tensor<4x3xf32>) -> tensor<4x3xf32>
+// CHECK-NEXT:   %1 = stablehlo.clamp %cst_0, %0, %cst : tensor<4x3xf32>
+// CHECK-NEXT:   return %1 : tensor<4x3xf32>
+// CHECK-NEXT: }