fix: minor fixes

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -15008,8 +15008,9 @@ struct IfBinaryOpToSelectBinaryOp final
                          Region &branchRegion, bool isTrueBranch,
                          PatternRewriter &rewriter) const {
     auto binOp = branchVal.getDefiningOp<BinaryOpType>();
-    if (!binOp)
+    if (!binOp) {
       return nullptr;
+    }
 
     Value base = nullptr;
     Value other = nullptr;
@@ -15026,26 +15027,31 @@ struct IfBinaryOpToSelectBinaryOp final
       matchLhs = false;
     }
 
-    if (!base || !other)
+    if (!base || !other) {
       return nullptr;
+    }
 
-    bool isCommutative = binOp->template hasTrait<mlir::OpTrait::IsCommutative>();
+    bool isCommutative =
+        binOp->template hasTrait<mlir::OpTrait::IsCommutative>();
 
     // For non-commutative ops, we only support hoisting if the identity works.
     // We assume getIdentityValueForOp returns a right identity.
     // Thus, for non-commutative ops, the base MUST be the LHS.
-    if (!isCommutative && !matchLhs)
+    if (!isCommutative && !matchLhs) {
       return nullptr;
+    }
 
     // Check if 'other' can be hoisted (not defined in the current branch)
-    if (&branchRegion == other.getParentRegion())
+    if (&branchRegion == other.getParentRegion()) {
       return nullptr;
+    }
 
     auto elemType = cast<ShapedType>(branchVal.getType()).getElementType();
     Value identity = stablehlo::getIdentityValueForOp<BinaryOpType>(
         rewriter, op.getLoc(), elemType);
-    if (!identity)
+    if (!identity) {
       return nullptr;
+    }
 
     // Broadcast identity to match the type of 'other'
     auto otherType = cast<RankedTensorType>(other.getType());
@@ -15060,10 +15066,11 @@ struct IfBinaryOpToSelectBinaryOp final
 
     Value selected = stablehlo::SelectOp::create(rewriter, op.getLoc(), pred,
                                                  trueSelectVal, falseSelectVal);
-    if (matchLhs)
+    if (matchLhs) {
       return BinaryOpType::create(rewriter, op.getLoc(), base, selected);
-    else
+    } else {
       return BinaryOpType::create(rewriter, op.getLoc(), selected, base);
+    }
   }
 };
 

test/lit_tests/if_binop_to_select.mlir

@@ -38,7 +38,7 @@ module {
   // Case with Mul and identity 1.0
   func.func @hoist_mul(%pred: tensor<i1>, %base: tensor<f32>, %other: tensor<f32>) -> tensor<f32> {
     %0 = "stablehlo.if"(%pred) ({
-      %1 = stablehlo.mul %other, %base : tensor<f32>
+      %1 = stablehlo.multiply %other, %base : tensor<f32>
       "stablehlo.return"(%1) : (tensor<f32>) -> ()
     }, {
       "stablehlo.return"(%base) : (tensor<f32>) -> ()
@@ -48,7 +48,7 @@ module {
   // CHECK-LABEL: func.func @hoist_mul
   // CHECK: %[[ONE:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
   // CHECK: %[[SEL:.+]] = stablehlo.select %{{.+}}, %{{.+}}, %[[ONE]]
-  // CHECK: %[[MUL:.+]] = stablehlo.mul %{{.+}}, %[[SEL]]
+  // CHECK: %[[MUL:.+]] = stablehlo.multiply %{{.+}}, %[[SEL]]
 
   // Case with Max and identity -Inf
   func.func @hoist_max(%pred: tensor<i1>, %base: tensor<f32>, %other: tensor<f32>) -> tensor<f32> {