feat: common syrk optimization patterns (#1653)

* feat: common syrk optimization patterns

* feat: syrk output is always symmetric

* fix: address review comment

* fix: docs

Author: avik-pal

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -427,7 +427,11 @@ def SyrkOp: EnzymeXLA_Op<"blas.syrk", [Pure, SameOperandsAndResultElementType]>
   let summary = "Multiplication involving a symmetric matrix";
 
   let description = [{
-    C := alpha*A*A^T + beta*C, or C := alpha*A^T*A + beta*C, where alpha and beta are scalars. C must be a n x n symmetric matrix."
+    C := alpha*A*A^T + beta*C, or C := alpha*A^T*A + beta*C, where alpha and beta are
+    scalars. C must be a n x n symmetric matrix.
+
+    If `fill` is present, then both the upper and lower triangles of the matrix are filled.
+    Otherwise the values in the non-uplo part of the matrix are undefined.
   }];
 
   let arguments = (ins
@@ -436,7 +440,8 @@ def SyrkOp: EnzymeXLA_Op<"blas.syrk", [Pure, SameOperandsAndResultElementType]>
     TensorFloat:$alpha,
     TensorFloat:$beta,
     EnzymeXLA_LapackUploAttr:$uplo,
-    DefaultValuedAttr<EnzymeXLA_LapackTransposeAttr, "::mlir::enzymexla::LapackTranspose::none">:$transpose
+    DefaultValuedAttr<EnzymeXLA_LapackTransposeAttr, "::mlir::enzymexla::LapackTranspose::none">:$transpose,
+    OptionalAttr<UnitAttr>:$fill
   );
 
   let results = (outs

src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp

@@ -3166,44 +3166,6 @@ struct IfOpEnzymeOpsRemover
   }
 };
 
-Value getScalarInitValue(Operation *op, OpBuilder &builder) {
-  if (!op)
-    return nullptr;
-
-  // Splatted Constant
-  SplatElementsAttr elems;
-  if (matchPattern(op, m_Constant(&elems))) {
-    auto scalarElemType = RankedTensorType::get(
-        {}, cast<TensorType>(op->getResult(0).getType()).getElementType());
-    auto constInit = ConstantOp::create(builder, op->getLoc(), scalarElemType,
-                                        elems.resizeSplat(scalarElemType));
-    return constInit;
-  }
-
-  // BroadcastInDim / Reshape
-  if (isa<stablehlo::BroadcastInDimOp, stablehlo::ReshapeOp>(op)) {
-    if (cast<RankedTensorType>(op->getOperand(0).getType()).getRank() == 0) {
-      return op->getOperand(0);
-    }
-  }
-
-  // Convert
-  if (auto convertOp = dyn_cast<stablehlo::ConvertOp>(op)) {
-    auto scalar =
-        getScalarInitValue(convertOp.getOperand().getDefiningOp(), builder);
-    if (scalar) {
-      auto convertOutElemType =
-          cast<RankedTensorType>(convertOp.getResult().getType())
-              .getElementType();
-      return stablehlo::ConvertOp::create(
-          builder, op->getLoc(), RankedTensorType::get({}, convertOutElemType),
-          scalar);
-    }
-  }
-
-  return nullptr;
-}
-
 struct SHLOReduceOpBatchInterface
     : public BatchOpInterface::ExternalModel<SHLOReduceOpBatchInterface,
                                              ReduceOp> {
@@ -3234,8 +3196,7 @@ struct SHLOReduceOpBatchInterface
     newReduceInits.reserve(reduceOp.getInitValues().size());
     for (auto opValue : reduceOp.getInitValues()) {
       auto batchedInit = mapper.lookup(opValue);
-      auto scalarInit =
-          getScalarInitValue(batchedInit.getDefiningOp(), builder);
+      auto scalarInit = getScalarValue(batchedInit.getDefiningOp(), builder);
       if (!scalarInit) {
         // TODO: we need to support broadcasting inits, or do we?
         src->emitError("Unsupported reduce init for batched reduce");
@@ -3316,8 +3277,7 @@ struct SHLOReduceWindowOpBatchInterface
     newReduceWindowInits.reserve(reduceWindowOp.getInitValues().size());
     for (auto opValue : reduceWindowOp.getInitValues()) {
       auto batchedInit = mapper.lookup(opValue);
-      auto scalarInit =
-          getScalarInitValue(batchedInit.getDefiningOp(), builder);
+      auto scalarInit = getScalarValue(batchedInit.getDefiningOp(), builder);
       if (!scalarInit) {
         src->emitError(
             "Unsupported reduce window init for batched reduce window");

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -25871,6 +25871,227 @@ struct BinaryNegatedOperandsSimplify
   }
 };
 
+// currently limited to non-batched dot_general
+struct DotGeneralToSyrk
+    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
+                                     DotGeneralToSyrk> {
+  using CheckedOpRewritePattern<stablehlo::DotGeneralOp,
+                                DotGeneralToSyrk>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
+                                    PatternRewriter &rewriter) const {
+    auto dotDims = op.getDotDimensionNumbers();
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+
+    if (dotDims.getLhsBatchingDimensions().size() != 0 ||
+        dotDims.getRhsBatchingDimensions().size() != 0) {
+      return failure();
+    }
+
+    if (dotDims.getLhsContractingDimensions().size() != 1 ||
+        dotDims.getRhsContractingDimensions().size() != 1) {
+      return failure();
+    }
+
+    // check that transpose dimensions are [1,0]
+    auto isTrueTranspose = [](stablehlo::TransposeOp tOp) -> bool {
+      auto perm = tOp.getPermutation();
+      return perm.size() == 2 && perm[0] == 1 && perm[1] == 0;
+    };
+
+    auto lhsContractingDim = dotDims.getLhsContractingDimensions()[0];
+    auto rhsContractingDim = dotDims.getRhsContractingDimensions()[0];
+
+    Value syrkInput;
+    enzymexla::LapackTranspose lapackTranspose;
+
+    if (lhs == rhs && lhsContractingDim == rhsContractingDim) {
+      syrkInput = lhs;
+      if (lhsContractingDim == 1) {
+        lapackTranspose = enzymexla::LapackTranspose::none;
+      } else {
+        lapackTranspose = enzymexla::LapackTranspose::transpose;
+      }
+    }
+
+    if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
+      if (isTrueTranspose(lhsT) && lhsT.getOperand() == rhs &&
+          lhsContractingDim == 1 - rhsContractingDim) {
+        syrkInput = rhs;
+        if (rhsContractingDim == 1) {
+          lapackTranspose = enzymexla::LapackTranspose::none;
+        } else {
+          lapackTranspose = enzymexla::LapackTranspose::transpose;
+        }
+      }
+    }
+
+    if (auto rhsT = rhs.getDefiningOp<stablehlo::TransposeOp>()) {
+      if (isTrueTranspose(rhsT) && rhsT.getOperand() == lhs &&
+          rhsContractingDim == 1 - lhsContractingDim) {
+        syrkInput = lhs;
+        if (lhsContractingDim == 0) {
+          lapackTranspose = enzymexla::LapackTranspose::transpose;
+        } else {
+          lapackTranspose = enzymexla::LapackTranspose::none;
+        }
+      }
+    }
+
+    if (!syrkInput)
+      return failure();
+
+    auto elemType =
+        cast<RankedTensorType>(syrkInput.getType()).getElementType();
+    auto alphaType = RankedTensorType::get({}, elemType);
+
+    auto syrkOp = enzymexla::SyrkOp::create(
+        rewriter, op.getLoc(), op.getResult().getType(), syrkInput,
+        stablehlo::ConstantOp::create(
+            rewriter, op.getLoc(), op.getType(),
+            cast<ElementsAttr>(makeAttr(op.getType(), 0))),
+        stablehlo::ConstantOp::create(
+            rewriter, op.getLoc(), alphaType,
+            cast<ElementsAttr>(makeAttr(alphaType, 1))),
+        stablehlo::ConstantOp::create(
+            rewriter, op.getLoc(), alphaType,
+            cast<ElementsAttr>(makeAttr(alphaType, 0))),
+        enzymexla::LapackUploAttr::get(op.getContext(),
+                                       enzymexla::LapackUplo::U),
+        enzymexla::LapackTransposeAttr::get(op.getContext(), lapackTranspose),
+        rewriter.getUnitAttr());
+    rewriter.replaceOp(op, syrkOp.getResult());
+    return success();
+  }
+};
+
+struct TransposeSyrkToSyrk
+    : public CheckedOpRewritePattern<enzymexla::SyrkOp, TransposeSyrkToSyrk> {
+  using CheckedOpRewritePattern<enzymexla::SyrkOp,
+                                TransposeSyrkToSyrk>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(enzymexla::SyrkOp op,
+                                    PatternRewriter &rewriter) const {
+    auto input = op.getA();
+    if (cast<RankedTensorType>(input.getType()).getRank() != 2)
+      return failure(); // support only rank 2 matrices for now
+
+    auto transposeOp = input.getDefiningOp<stablehlo::TransposeOp>();
+    if (!transposeOp)
+      return failure();
+
+    auto perm = transposeOp.getPermutation();
+    if (perm.size() != 2 || perm[0] != 1 || perm[1] != 0)
+      return failure();
+
+    enzymexla::LapackTranspose lapackTranspose;
+    switch (op.getTranspose()) {
+    case enzymexla::LapackTranspose::none:
+      lapackTranspose = enzymexla::LapackTranspose::transpose;
+      break;
+    default:
+      lapackTranspose = enzymexla::LapackTranspose::none;
+    }
+
+    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
+        op, op.getResult().getType(), transposeOp.getOperand(), op.getC(),
+        op.getAlpha(), op.getBeta(), op.getUploAttr(),
+        enzymexla::LapackTransposeAttr::get(op.getContext(), lapackTranspose),
+        op.getFillAttr());
+    return success();
+  }
+};
+
+struct FuseMulIntoSyrk
+    : public CheckedOpRewritePattern<stablehlo::MulOp, FuseMulIntoSyrk> {
+  using CheckedOpRewritePattern<stablehlo::MulOp,
+                                FuseMulIntoSyrk>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::MulOp op,
+                                    PatternRewriter &rewriter) const {
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+
+    enzymexla::SyrkOp syrkOp;
+    Value other;
+
+    if (auto lhsSyrk = lhs.getDefiningOp<enzymexla::SyrkOp>()) {
+      syrkOp = lhsSyrk;
+      other = rhs;
+    } else if (auto rhsSyrk = rhs.getDefiningOp<enzymexla::SyrkOp>()) {
+      syrkOp = rhsSyrk;
+      other = lhs;
+    } else {
+      return failure();
+    }
+
+    Value scalarVal =
+        stablehlo::getScalarValue(other.getDefiningOp(), rewriter);
+    if (!scalarVal)
+      return failure();
+
+    auto newBeta = stablehlo::MulOp::create(rewriter, op.getLoc(),
+                                            syrkOp.getBeta(), scalarVal);
+    auto newAlpha = stablehlo::MulOp::create(rewriter, op.getLoc(),
+                                             syrkOp.getAlpha(), scalarVal);
+
+    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
+        op, syrkOp.getType(), syrkOp.getA(), syrkOp.getC(), newAlpha, newBeta,
+        syrkOp.getUploAttr(), syrkOp.getTransposeAttr(), syrkOp.getFillAttr());
+    return success();
+  }
+};
+
+struct FuseAddIntoSyrk
+    : public CheckedOpRewritePattern<stablehlo::AddOp,
+                                     FuseAddIntoSyrk>::CheckedOpRewritePattern {
+  using CheckedOpRewritePattern<stablehlo::AddOp,
+                                FuseAddIntoSyrk>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::AddOp op,
+                                    PatternRewriter &rewriter) const {
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+
+    enzymexla::SyrkOp syrkOp;
+    Value other;
+
+    if (auto lhsSyrk = lhs.getDefiningOp<enzymexla::SyrkOp>()) {
+      syrkOp = lhsSyrk;
+      other = rhs;
+    } else if (auto rhsSyrk = rhs.getDefiningOp<enzymexla::SyrkOp>()) {
+      syrkOp = rhsSyrk;
+      other = lhs;
+    } else {
+      return failure();
+    }
+
+    // we can fuse this addition iff the other operand is a symmetric matrix
+    if (!canApplySymmetricPattern(other, rewriter))
+      return failure();
+
+    auto oldBeta = syrkOp.getBeta();
+    auto bcastedBeta = stablehlo::BroadcastInDimOp::create(
+        rewriter, op.getLoc(), syrkOp.getType(), oldBeta,
+        rewriter.getDenseI64ArrayAttr({}));
+
+    auto scaledC = stablehlo::MulOp::create(rewriter, op.getLoc(),
+                                            syrkOp.getC(), bcastedBeta);
+
+    auto newC = stablehlo::AddOp::create(rewriter, op.getLoc(), scaledC, other);
+
+    auto newBeta = stablehlo::ConstantOp::create(
+        rewriter, op.getLoc(), oldBeta.getType(),
+        cast<ElementsAttr>(makeAttr(oldBeta.getType(), 1)));
+
+    rewriter.replaceOpWithNewOp<enzymexla::SyrkOp>(
+        op, syrkOp.getType(), syrkOp.getA(), newC, syrkOp.getAlpha(), newBeta,
+        syrkOp.getUploAttr(), syrkOp.getTransposeAttr(), syrkOp.getFillAttr());
+    return success();
+  }
+};
+
 ///////////////  End Imported from stablehlo
 
 // clang-format off
@@ -26423,6 +26644,12 @@ struct EnzymeHLOOptPass
     patterns.add<TransposeSymmetricSimplify>(context);
     patterns.add<FactorScalarsInDotGeneral>(context);
 
+    // syrk patterns
+    // currently disabled since lowering is missing
+    // patterns.add<DotGeneralToSyrk>(context);
+    patterns.add<TransposeSyrkToSyrk, FuseMulIntoSyrk, FuseAddIntoSyrk>(
+        context);
+
     // clang-format off
     patterns.add<
         WhileRepeatedInductionReduction,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2519,3 +2519,23 @@ def ApplyMultiplyNegatedOperandsSimplifyPatterns : EnzymeHLOPatternOp<
     "multiply_negated_operands_simplify"> {
   let patterns = ["BinaryNegatedOperandsSimplify<stablehlo::MulOp>"];
 }
+
+def ApplyDotGeneralToSyrkPatterns : EnzymeHLOPatternOp<
+    "dot_general_to_syrk"> {
+  let patterns = ["DotGeneralToSyrk"];
+}
+
+def ApplyTransposeSyrkToSyrkPatterns : EnzymeHLOPatternOp<
+    "transpose_syrk_to_syrk"> {
+  let patterns = ["TransposeSyrkToSyrk"];
+}
+
+def ApplyFuseMulIntoSyrkPatterns : EnzymeHLOPatternOp<
+    "fuse_mul_into_syrk"> {
+  let patterns = ["FuseMulIntoSyrk"];
+}
+
+def ApplyFuseAddIntoSyrkPatterns : EnzymeHLOPatternOp<
+    "fuse_add_into_syrk"> {
+  let patterns = ["FuseAddIntoSyrk"];
+}