Add partial symmetry detection

src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h

@@ -0,0 +1,114 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+namespace enzyme {
+
+// Represents the partial symmetry of a tensor as a partition of its dimensions.
+class PartialSymmetryAnnotation {
+public:
+  PartialSymmetryAnnotation() : dimensionSetIDs() {}
+
+  explicit PartialSymmetryAnnotation(ArrayRef<int64_t> dimensionSetIDs);
+
+  static PartialSymmetryAnnotation createUninitialized(int64_t rank);
+  static PartialSymmetryAnnotation createNotSymmetric(int64_t rank);
+  static PartialSymmetryAnnotation createFullySymmetric(int64_t rank);
+
+  bool isSymmetric(int64_t i, int64_t j) const;
+  int64_t getSetId(int64_t i) const { return dimensionSetIDs[i]; }
+  int64_t getRank() const { return dimensionSetIDs.size(); }
+
+  static PartialSymmetryAnnotation meet(const PartialSymmetryAnnotation &lhs,
+                                        const PartialSymmetryAnnotation &rhs);
+  static PartialSymmetryAnnotation join(const PartialSymmetryAnnotation &lhs,
+                                        const PartialSymmetryAnnotation &rhs);
+
+  static PartialSymmetryAnnotation
+  propagateTranspose(const PartialSymmetryAnnotation &annotation,
+                     ArrayRef<int64_t> permutation);
+
+  static PartialSymmetryAnnotation
+  propagateBroadcastInDim(const PartialSymmetryAnnotation &annotation,
+                          int64_t outputRank,
+                          ArrayRef<int64_t> broadcastDimensions);
+
+  static PartialSymmetryAnnotation
+  propagateDotGeneral(const PartialSymmetryAnnotation &lhsAnnotation,
+                      const PartialSymmetryAnnotation &rhsAnnotation,
+                      int64_t resultRank, ArrayRef<int64_t> lhsBatchingDims,
+                      ArrayRef<int64_t> rhsBatchingDims,
+                      ArrayRef<int64_t> lhsContractingDims,
+                      ArrayRef<int64_t> rhsContractingDims, bool rhsAliasesLhs,
+                      ArrayRef<int64_t> rhsDimToLhs);
+
+  static PartialSymmetryAnnotation checkConstant(DenseElementsAttr attr);
+
+  static PartialSymmetryAnnotation
+  propagateElementwiseBinary(const PartialSymmetryAnnotation &lhsAnnotation,
+                             const PartialSymmetryAnnotation &rhsAnnotation,
+                             int64_t resultRank, bool rhsAliasesLhs,
+                             ArrayRef<int64_t> rhsDimToLhs);
+
+  bool operator==(const PartialSymmetryAnnotation &other) const {
+    return dimensionSetIDs == other.dimensionSetIDs;
+  }
+
+  SmallVector<SmallVector<int64_t>> getDimensionSets() const;
+  static PartialSymmetryAnnotation
+  createFromDimensionSets(int64_t rank,
+                          ArrayRef<ArrayRef<int64_t>> dimensionSets);
+  static std::optional<PartialSymmetryAnnotation> createFromIR(Value val);
+
+  void print(raw_ostream &os) const;
+
+private:
+  SmallVector<int64_t> dimensionSetIDs;
+
+  void canonicalize();
+  void uniteDimensionSets(int64_t rank, int64_t i, int64_t j);
+};
+
+class PartialSymmetryLattice : public dataflow::AbstractSparseLattice {
+public:
+  using AbstractSparseLattice::AbstractSparseLattice;
+
+  PartialSymmetryLattice(Value v);
+
+  ChangeResult meet(const AbstractSparseLattice &rhs) override;
+  ChangeResult meet(const PartialSymmetryLattice &rhs);
+
+  void print(raw_ostream &os) const override;
+
+  const PartialSymmetryAnnotation &getValue() const { return value; }
+  void setValue(const PartialSymmetryAnnotation &v) { value = v; }
+
+private:
+  PartialSymmetryAnnotation value;
+};
+
+class PartialSymmetryAnalysis
+    : public dataflow::SparseForwardDataFlowAnalysis<PartialSymmetryLattice> {
+public:
+  using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+
+  void setToEntryState(PartialSymmetryLattice *lattice) override;
+
+  LogicalResult
+  visitOperation(Operation *op,
+                 ArrayRef<const PartialSymmetryLattice *> operands,
+                 ArrayRef<PartialSymmetryLattice *> results) override;
+};
+
+} // namespace enzyme
+} // namespace mlir

src/enzyme_ad/jax/BUILD

@@ -840,6 +840,7 @@ cc_library(
 cc_library(
     name = "XLADerivatives",
     srcs = glob([
+        "Analysis/*.cpp",
         "Implementations/*.cpp",
         "Passes/*.cpp",
         "Dialect/*.cpp",
@@ -849,6 +850,7 @@ cc_library(
         "Utils.cpp",
     ],
     hdrs = glob([
+        "Analysis/*.h",
         "Implementations/*.h",
         "Passes/*.h",
         "Dialect/*.h",

src/enzyme_ad/jax/Dialect/EnzymeXLAAttrs.td

@@ -133,4 +133,36 @@ def EnzymeXLA_GuaranteedAnalysisResult : I32EnumAttr<"GuaranteedAnalysisResult",
 def EnzymeXLA_GuaranteedAnalysisResultAttr : EnumAttr<EnzymeXLA_Dialect,
     EnzymeXLA_GuaranteedAnalysisResult, "guaranteed">;
 
+def EnzymeXLA_SymmetricDimensionSetAttr : AttrDef<EnzymeXLA_Dialect,
+    "SymmetricDimensionSet"> {
+  let summary = "A set of symmetric dimension indices";
+  let cppNamespace = "::mlir::enzymexla";
+
+  let parameters = (ins
+    "DenseI64ArrayAttr":$dimensions
+  );
+
+  let assemblyFormat = [{
+    `<` $dimensions `>`
+  }];
+
+  let mnemonic = "symmetric_dimension_set";
+}
+
+def EnzymeXLA_PartialSymmetryAnalysisResultAttr : AttrDef<EnzymeXLA_Dialect,
+    "PartialSymmetryAnalysisResult"> {
+  let summary = "Sets of partially symmetric dimensions";
+  let cppNamespace = "::mlir::enzymexla";
+
+  let parameters = (ins
+    ArrayRefParameter<"SymmetricDimensionSetAttr">:$values
+  );
+
+  let assemblyFormat = [{
+    `<` $values `>`
+  }];
+
+  let mnemonic = "partial_symmetry";
+}
+
 #endif // ENZYMEXLA_ATTRS

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -24,9 +24,11 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Visitors.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "shardy/dialect/sdy/ir/utils.h"
+#include "src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h"
 #include "src/enzyme_ad/jax/CheckedRewrite.h"
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
@@ -55,6 +57,7 @@
 #include "llvm/ADT/MapVector.h"
 #include <iterator>
 #include <numeric>
+#include <optional>
 #define DEBUG_TYPE "enzymehloopt"
 
 namespace mlir {
@@ -6957,6 +6960,42 @@ struct TransposeSymmetricSimplify
   }
 };
 
+struct TransposePartialSymmetrySimplify
+    : public CheckedOpRewritePattern<stablehlo::TransposeOp,
+                                     TransposePartialSymmetrySimplify> {
+  using CheckedOpRewritePattern<
+      stablehlo::TransposeOp,
+      TransposePartialSymmetrySimplify>::CheckedOpRewritePattern;
+
+  LogicalResult matchAndRewriteImpl(stablehlo::TransposeOp op,
+                                    PatternRewriter &rewriter) const {
+    auto operand = op.getOperand();
+    auto perm = op.getPermutation();
+
+    auto annotationOpt =
+        enzyme::PartialSymmetryAnnotation::createFromIR(operand);
+    if (!annotationOpt.has_value())
+      return failure();
+
+    auto annotation = annotationOpt.value();
+
+    bool isIdentity = true;
+    for (int64_t i = 0; i < (int64_t)perm.size(); ++i) {
+      if (annotation.getSetId(i) != annotation.getSetId(perm[i])) {
+        isIdentity = false;
+        break;
+      }
+    }
+
+    if (isIdentity) {
+      rewriter.replaceOp(op, operand);
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 struct NoNanSelfSubSimplify
     : public NoNanCheckedOpRewritePattern<stablehlo::SubtractOp,
                                           NoNanSelfSubSimplify> {
@@ -26653,7 +26692,8 @@ struct EnzymeHLOOptPass
                  NoNanAddSubSimplify, NoNanMulSimplify, NoNanDivSimplify>(
         (no_nan || all_finite), context);
 
-    patterns.add<TransposeSymmetricSimplify>(context);
+    patterns.add<TransposeSymmetricSimplify, TransposePartialSymmetrySimplify>(
+        context);
     patterns.add<FactorScalarsInDotGeneral>(context);
 
     // syrk patterns

src/enzyme_ad/jax/Passes/PartialSymmetryAnnotate.cpp

@@ -0,0 +1,99 @@
+#include "src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h"
+#include "src/enzyme_ad/jax/Passes/Passes.h"
+
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "src/enzyme_ad/jax/Dialect/Dialect.h"
+#include "src/enzyme_ad/jax/Dialect/Ops.h"
+#include "stablehlo/dialect/StablehloOps.h"
+
+#define DEBUG_TYPE "partial-symmetry-annotate"
+
+namespace mlir {
+namespace enzyme {
+#define GEN_PASS_DEF_PARTIALSYMMETRYANNOTATEPASS
+#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
+} // namespace enzyme
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::dataflow;
+using namespace mlir::enzyme;
+
+namespace {
+
+class PartialSymmetryAnnotatePass
+    : public enzyme::impl::PartialSymmetryAnnotatePassBase<
+          PartialSymmetryAnnotatePass> {
+public:
+  using Base::Base;
+
+  void runOnOperation() override {
+    DataFlowSolver solver;
+
+    solver.load<enzyme::PartialSymmetryAnalysis>();
+    solver.load<dataflow::DeadCodeAnalysis>();
+    solver.load<dataflow::SparseConstantPropagation>();
+
+    if (failed(solver.initializeAndRun(getOperation()))) {
+      return signalPassFailure();
+    }
+
+    auto mod = getOperation();
+
+    // Annotate all operations with partial symmetry information
+    mod->walk([&](Operation *op) {
+      SmallVector<Attribute> partialSymmetryAttrs;
+      bool anyKnown = false;
+
+      for (auto result : op->getResults()) {
+        auto *state =
+            solver.lookupState<enzyme::PartialSymmetryLattice>(result);
+        if (!state) {
+          continue;
+        }
+
+        auto dimensionSets = state->getValue().getDimensionSets();
+
+        SmallVector<enzymexla::SymmetricDimensionSetAttr> dimensionSetAttrs;
+        for (const auto &set : dimensionSets) {
+          if (set.size() > 1) {
+            anyKnown = true;
+            auto denseAttr = DenseI64ArrayAttr::get(mod.getContext(), set);
+            auto dimensionSetAttr = enzymexla::SymmetricDimensionSetAttr::get(
+                mod.getContext(), denseAttr);
+            dimensionSetAttrs.push_back(dimensionSetAttr);
+          }
+        }
+
+        if (dimensionSetAttrs.empty()) {
+          continue;
+        }
+
+        auto partialSymmetry =
+            enzymexla::PartialSymmetryAnalysisResultAttr::get(
+                mod.getContext(), dimensionSetAttrs);
+        partialSymmetryAttrs.push_back(partialSymmetry);
+      }
+
+      if (anyKnown) {
+        op->setAttr("enzymexla.partial_symmetry",
+                    ArrayAttr::get(mod.getContext(), partialSymmetryAttrs));
+      }
+
+      return WalkResult::advance();
+    });
+  }
+};
+
+} // namespace

src/enzyme_ad/jax/Passes/Passes.td

@@ -1083,6 +1083,15 @@ def EnzymeBatchToStableHLOPass : Pass<"enzyme-batch-to-stablehlo"> {
   ];
 }
 
+def PartialSymmetryAnnotatePass : Pass<"partial-symmetry-annotate", "ModuleOp"> {
+  let summary = "Annotate operations using partial symmetry analysis";
+  let dependentDialects = [
+    "stablehlo::StablehloDialect",
+    "enzymexla::EnzymeXLADialect",
+    "func::FuncDialect",
+  ];
+}
+
 def ConvertAllConstantsToSplattedConstantPass : Pass<"convert-all-constants-to-splatted-constant", "ModuleOp"> {
   let summary = "Convert all constants to splatted constants. This is supposed to be used for debugging purposes when dumping the module.";
   let dependentDialects = [

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -608,6 +608,11 @@ def ApplyTransposeSymmetricSimplify : EnzymeHLOPatternOp<
   let patterns = ["TransposeSymmetricSimplify"];
 }
 
+def ApplyTransposePartialSymmetrySimplify : EnzymeHLOPatternOp<
+    "transpose_partial_symmetry_simplify"> {
+  let patterns = ["TransposePartialSymmetrySimplify"];
+}
+
 def ApplyFactorScalarsInDotGeneral : EnzymeHLOPatternOp<
     "factor_scalars_in_dot_general"> {
   let patterns = ["FactorScalarsInDotGeneral"];

src/enzyme_ad/jax/Utils.h

@@ -319,6 +319,9 @@ bool mayAlias(mlir::MemoryEffects::EffectInstance a,
 
 bool mayAlias(mlir::MemoryEffects::EffectInstance a, mlir::Value b);
 
+bool checkNotEqual(llvm::APInt a, llvm::APInt b);
+bool checkNotEqual(llvm::APFloat a, llvm::APFloat b);
+
 bool canApplyNoNanPattern(bool allowOnFloatingPointMath, Type Ty);
 bool canApplyNoNanPattern(bool allowOnFloatingPointMath, Type Ty,
                           mlir::Operation *op, PatternRewriter &rewriter);