feat: simplify diagonal accesses for a dot_general op (#1614)

* feat: simplify diagonal accesses for a dot_general op

* test: update

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -25483,6 +25483,145 @@ struct DynamicSliceSimplify
   }
 };
 
+// TODO: generalize to higher ranked tensors
+// TODO: if we determine that all accesses are on some offset diagonal,
+// we can still replace it will a multiply combined with pad/slice
+// If we prove that only the diagonal elements of a dot_general are accessed,
+// we replace the dot_general with a cheaper multiply op. Note that
+// this implies `diag(new_op(A, B)) = diag(A x B)` however
+// `new_op(A, B) != A x B`
+struct DotGeneralOnlyDiagonalAccess
+    : public CheckedOpRewritePattern<stablehlo::DotGeneralOp,
+                                     DotGeneralOnlyDiagonalAccess> {
+  using CheckedOpRewritePattern<
+      stablehlo::DotGeneralOp,
+      DotGeneralOnlyDiagonalAccess>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::DotGeneralOp op,
+                                    PatternRewriter &rewriter) const {
+    auto resTy = cast<RankedTensorType>(op.getType());
+    if (resTy.getRank() != 2)
+      return failure();
+
+    auto M = resTy.getDimSize(0);
+    auto N = resTy.getDimSize(1);
+    auto diagLen = std::min(M, N);
+
+    auto lhs = op.getLhs();
+    auto rhs = op.getRhs();
+    auto dotDimNumbers = op.getDotDimensionNumbers();
+    auto lhsContractingDims = dotDimNumbers.getLhsContractingDimensions();
+    auto rhsContractingDims = dotDimNumbers.getRhsContractingDimensions();
+    auto lhsBatchingDims = dotDimNumbers.getLhsBatchingDimensions();
+    auto rhsBatchingDims = dotDimNumbers.getRhsBatchingDimensions();
+
+    if (lhsContractingDims.size() != 1 || rhsContractingDims.size() != 1 ||
+        lhsBatchingDims.size() != 0 || rhsBatchingDims.size() != 0)
+      return failure();
+
+    llvm::SetVector<Operation *> opsToReplace;
+    llvm::SmallPtrSet<Operation *, 4> seenOps;
+    for (auto user : op->getUsers()) {
+      if (seenOps.count(user))
+        continue;
+      if (!enzyme::allAccessesAreOnMainDiagonal(user, opsToReplace))
+        return failure();
+      seenOps.insert(user);
+    }
+
+    if (opsToReplace.empty())
+      return failure();
+
+    // rewrite the dot_general to a multiply.
+    // we insert transpose ops here, but those will get removed later
+    auto lhsContractDim = lhsContractingDims[0];
+    auto rhsContractDim = rhsContractingDims[0];
+    // result[i, i] = sum_k (lhs[i, k] * rhs[k, i])
+    //              = reduce_sum(lhs[i, :] * rhs[:, i])
+    auto lhsNonContractDim = 1 - lhsContractDim;
+    auto rhsNonContractDim = 1 - rhsContractDim;
+
+    if (lhsContractDim == 0) {
+      // move to dim = 1
+      lhs = stablehlo::TransposeOp::create(
+          rewriter, op.getLoc(), lhs, rewriter.getDenseI64ArrayAttr({1, 0}));
+    }
+    lhs = stablehlo::SliceOp::create(
+        rewriter, op.getLoc(), lhs, rewriter.getDenseI64ArrayAttr({0, 0}),
+        rewriter.getDenseI64ArrayAttr(
+            {diagLen, cast<ShapedType>(lhs.getType()).getDimSize(1)}),
+        rewriter.getDenseI64ArrayAttr({1, 1})); // [DiagSize, C]
+
+    if (rhsContractDim == 0) {
+      // move to dim = 1
+      rhs = stablehlo::TransposeOp::create(
+          rewriter, op.getLoc(), rhs, rewriter.getDenseI64ArrayAttr({1, 0}));
+    }
+    rhs = stablehlo::SliceOp::create(
+        rewriter, op.getLoc(), rhs, rewriter.getDenseI64ArrayAttr({0, 0}),
+        rewriter.getDenseI64ArrayAttr(
+            {diagLen, cast<ShapedType>(rhs.getType()).getDimSize(1)}),
+        rewriter.getDenseI64ArrayAttr({1, 1})); // [DiagSize, C]
+
+    auto newMul = stablehlo::MulOp::create(rewriter, op.getLoc(), lhs,
+                                           rhs); // [DiagSize, C]
+
+    auto elemTy = cast<RankedTensorType>(newMul.getType()).getElementType();
+    auto tenElemTy = RankedTensorType::get({}, elemTy);
+    auto reduceOp = stablehlo::ReduceOp::create(
+        rewriter, op.getLoc(), ValueRange(newMul.getResult()),
+        ValueRange(stablehlo::ConstantOp::create(
+                       rewriter, op.getLoc(), tenElemTy,
+                       cast<ElementsAttr>(makeAttr(tenElemTy, 0)))
+                       .getResult()),
+        {1});
+
+    {
+      Region &region = reduceOp.getBody();
+      Block *block = rewriter.createBlock(&region);
+      block->addArgument(tenElemTy, op.getLoc());
+      block->addArgument(tenElemTy, op.getLoc());
+
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointToStart(block);
+      auto addOp = stablehlo::AddOp::create(
+          rewriter, op.getLoc(), block->getArgument(0), block->getArgument(1));
+      stablehlo::ReturnOp::create(rewriter, op.getLoc(), addOp.getResult());
+    }
+
+    for (auto &opToReplace : opsToReplace) {
+      if (auto sliceOp = dyn_cast<stablehlo::SliceOp>(opToReplace)) {
+        replaceSliceOp(rewriter, sliceOp, reduceOp, M, N, diagLen);
+      } else {
+        assert(false && "Unknown op to replace. open an issue on github");
+      }
+    }
+
+    return success();
+  }
+
+private:
+  void replaceSliceOp(PatternRewriter &rewriter, stablehlo::SliceOp sliceOp,
+                      stablehlo::ReduceOp reduceOp, int64_t M, int64_t N,
+                      int64_t diagLen) const {
+    int64_t start = sliceOp.getStartIndices()[0];
+    int64_t limit = sliceOp.getLimitIndices()[0];
+    int64_t stride = sliceOp.getStrides()[0];
+    int64_t diagStride = N + 1;
+
+    int64_t newStart = start / diagStride;
+    int64_t newLimit = (limit - 1) / diagStride + 1;
+    int64_t newStride = stride / diagStride;
+
+    rewriter.setInsertionPoint(sliceOp);
+    rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(
+        sliceOp, reduceOp.getResult(0),
+        rewriter.getDenseI64ArrayAttr({newStart}),
+        rewriter.getDenseI64ArrayAttr({newLimit}),
+        rewriter.getDenseI64ArrayAttr({newStride}));
+  }
+};
+
 ///////////////  End Imported from stablehlo
 
 // clang-format off
@@ -26117,7 +26256,8 @@ struct EnzymeHLOOptPass
         RemoveNoOpsFromWhileLoop,
         WhileIsCopySimplify,
         SplitVariadicScatterOp,
-        DynamicSliceSimplify
+        DynamicSliceSimplify,
+        DotGeneralOnlyDiagonalAccess
       >(context);
 
     patterns.add<

src/enzyme_ad/jax/Passes/StructuredTensors.cpp

@@ -5,6 +5,8 @@
 #include "src/enzyme_ad/jax/Utils.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
+#include "llvm/ADT/SetVector.h"
+
 namespace mlir {
 namespace enzyme {
 
@@ -256,5 +258,89 @@ std::optional<IotaLikeTensor> detectIotaLikeTensor(mlir::Value tensor) {
   return result;
 }
 
+bool allAccessesAreOnMainDiagonalPostReshape(stablehlo::ReshapeOp op,
+                                             stablehlo::SliceOp sliceOp) {
+  auto reshapeInTy = cast<RankedTensorType>(op.getOperand().getType());
+  auto reshapeOutTy = cast<RankedTensorType>(op.getType());
+
+  if (reshapeOutTy.getRank() != 1 ||
+      reshapeInTy.getRank() != 2) // [M, N] -> [M * N] vector
+    return false;
+
+  auto M = reshapeInTy.getDimSize(0);
+  auto N = reshapeInTy.getDimSize(1);
+  auto diagLen = std::min(M, N);
+  auto diagStride = N + 1;
+
+  int64_t start = sliceOp.getStartIndices()[0];
+  int64_t limit = sliceOp.getLimitIndices()[0];
+  int64_t stride = sliceOp.getStrides()[0];
+
+  if (stride % diagStride != 0)
+    return false;
+
+  // start can be on any of the diagonal elements
+  if (start % diagStride != 0)
+    return false;
+
+  if (limit > M * N)
+    return false; // technically this is illegal
+
+  // sanity check
+  int64_t count = (limit - start + stride - 1) / stride;
+  if (count <= 0 || count > diagLen)
+    return false;
+
+  return true;
+}
+
+bool allAccessesAreOnMainDiagonalPostReshape(
+    stablehlo::ReshapeOp op, Operation *user,
+    llvm::SetVector<Operation *> &opsToReplace) {
+  if (auto sliceOp = dyn_cast<stablehlo::SliceOp>(user)) {
+    if (allAccessesAreOnMainDiagonalPostReshape(op, sliceOp)) {
+      opsToReplace.insert(sliceOp);
+      return true;
+    }
+    return false;
+  }
+  return false;
+}
+
+bool allAccessesAreOnMainDiagonal(Operation *op,
+                                  llvm::SetVector<Operation *> &opsToReplace) {
+  if (auto reshapeOp = dyn_cast<stablehlo::ReshapeOp>(op)) {
+    return allAccessesAreOnMainDiagonal(reshapeOp, opsToReplace);
+  } else if (auto gatherOp = dyn_cast<stablehlo::GatherOp>(op)) {
+    return allAccessesAreOnMainDiagonal(gatherOp, opsToReplace);
+  }
+  return false;
+}
+
+bool allAccessesAreOnMainDiagonal(stablehlo::ReshapeOp op,
+                                  llvm::SetVector<Operation *> &opsToReplace) {
+  auto reshapeInTy = cast<RankedTensorType>(op.getOperand().getType());
+  if (reshapeInTy.getRank() != 2) // [M, N] matrix
+    return false;                 // quick exit
+
+  llvm::SmallPtrSet<Operation *, 4> seenOps;
+  for (auto user : op->getUsers()) {
+    if (seenOps.count(user))
+      continue;
+
+    if (!allAccessesAreOnMainDiagonalPostReshape(op, user, opsToReplace))
+      return false;
+
+    seenOps.insert(user);
+  }
+
+  return true;
+}
+
+bool allAccessesAreOnMainDiagonal(stablehlo::GatherOp op,
+                                  llvm::SetVector<Operation *> &opsToReplace) {
+  return false; // TODO: implement this where we are doing gather with iota
+}
+
 } // namespace enzyme
 } // namespace mlir

src/enzyme_ad/jax/Passes/StructuredTensors.h

@@ -5,6 +5,8 @@
 #include "src/enzyme_ad/jax/Utils.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
+#include "llvm/ADT/SetVector.h"
+
 #include <optional>
 
 namespace mlir {
@@ -27,5 +29,15 @@ struct IotaLikeTensor {
 
 std::optional<IotaLikeTensor> detectIotaLikeTensor(mlir::Value tensor);
 
+// TODO: we can do a full analysis and return if the access is on a specific set
+// of diagonals. Checks that all accesses for this Op and its users thereoff are
+// along the diagonal.
+bool allAccessesAreOnMainDiagonal(
+    mlir::Operation *op, llvm::SetVector<mlir::Operation *> &opsToReplace);
+bool allAccessesAreOnMainDiagonal(
+    stablehlo::ReshapeOp op, llvm::SetVector<mlir::Operation *> &opsToReplace);
+bool allAccessesAreOnMainDiagonal(
+    stablehlo::GatherOp op, llvm::SetVector<mlir::Operation *> &opsToReplace);
+
 } // namespace enzyme
 } // namespace mlir

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -2472,3 +2472,8 @@ def EnzymeHLOUnroll : EnzymeHLOParameterizedPatternOp<
     }
   }];
 }
+
+def ApplyDotGeneralOnlyDiagonalAccessPatterns : EnzymeHLOPatternOp<
+  "dot_general_only_diagonal_access"> {
+  let patterns = ["DotGeneralOnlyDiagonalAccess"];
+}

src/enzyme_ad/jax/primitives.py

@@ -346,6 +346,7 @@ def optimization_passes(
         "dynamic_pad_to_pad",
         "remove_no_ops_from_while_loop",
         "while_is_copy_simplify",
+        "dot_general_only_diagonal_access",
     ]
 
     # constant propagation patterns