Add n-dim transpose removal opt

Author: gaurav-arya

src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.cpp

@@ -386,6 +386,25 @@ PartialSymmetryAnnotation::getDimensionSets() const {
   return result;
 }
 
+PartialSymmetryAnnotation
+PartialSymmetryAnnotation::fromDimensionSets(int64_t rank,
+                                             ArrayRef<ArrayRef<int64_t>> dimensionSets) {
+  SmallVector<int64_t> dimensionSetIDs(rank);
+  for (int64_t i = 0; i < rank; ++i) {
+    dimensionSetIDs[i] = i;
+  }
+
+  // Note that dimensionSets is not assumed to be a complete partition.
+  // Missing dimensions are treated as separate sets.
+  for (auto dims : dimensionSets) {
+    for (int64_t i = 1; i < (int64_t)dims.size(); ++i) {
+      dimensionSetIDs[dims[i]] = dimensionSetIDs[dims[0]];
+    }
+  }
+
+  return PartialSymmetryAnnotation(dimensionSetIDs);
+}
+
 void PartialSymmetryAnnotation::print(raw_ostream &os) const {
   auto dimensionSets = getDimensionSets();
   os << "{";

src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h

@@ -68,6 +68,9 @@ class PartialSymmetryAnnotation {
 
   SmallVector<SmallVector<int64_t>> getDimensionSets() const;
 
+  static PartialSymmetryAnnotation
+  fromDimensionSets(int64_t rank, ArrayRef<ArrayRef<int64_t>> dimensionSets);
+
   void print(raw_ostream &os) const;
 
 private:

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -33,6 +33,7 @@
 #include "src/enzyme_ad/jax/Implementations/WhileLoopInfo.h"
 #include "src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "src/enzyme_ad/jax/Analysis/PartialSymmetryAnalysis.h"
 #include "src/enzyme_ad/jax/Passes/StructuredTensors.h"
 #include "src/enzyme_ad/jax/Utils.h"
 #include "stablehlo/dialect/Base.h"
@@ -55,6 +56,7 @@
 #include "llvm/ADT/MapVector.h"
 #include <iterator>
 #include <numeric>
+#include <optional>
 #define DEBUG_TYPE "enzymehloopt"
 
 namespace mlir {
@@ -6953,6 +6955,82 @@ struct TransposeSymmetricSimplify
   }
 };
 
+static std::optional<enzyme::PartialSymmetryAnnotation>
+getPartialSymmetryFromAttr(Value val) {
+  auto op = val.getDefiningOp();
+  if (!op)
+    return std::nullopt;
+
+  auto arrayAttr =
+      op->getAttrOfType<ArrayAttr>("enzymexla.partial_symmetry");
+  if (!arrayAttr || arrayAttr.empty())
+    return std::nullopt;
+
+  // Get the result number for this value
+  auto opResult = dyn_cast<OpResult>(val);
+  if (!opResult)
+    return std::nullopt;
+
+  auto resultNumber = opResult.getResultNumber();
+  if (resultNumber >= arrayAttr.size())
+    return std::nullopt;
+
+  auto partialSymmetryAttr = dyn_cast<enzymexla::PartialSymmetryAnalysisResultAttr>(
+      arrayAttr[resultNumber]);
+  if (!partialSymmetryAttr)
+    return std::nullopt;
+
+  auto dimensionSetAttrs = partialSymmetryAttr.getValues();
+  auto rank = cast<RankedTensorType>(val.getType()).getRank();
+
+  SmallVector<ArrayRef<int64_t>> dimensionSets;
+  for (auto dimensionSetAttr : dimensionSetAttrs) {
+    auto dims = dimensionSetAttr.getDimensions().asArrayRef();
+    dimensionSets.push_back(dims);
+  }
+
+  return enzyme::PartialSymmetryAnnotation::fromDimensionSets(rank, dimensionSets);
+}
+
+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();
+
+    // Get partial symmetry annotation from the operand
+    auto annotationOpt = getPartialSymmetryFromAttr(operand);
+    if (!annotationOpt.has_value())
+      return failure();
+    
+    auto annotation = annotationOpt.value();
+
+    // Check if the transpose is an identity based on partial symmetry
+    // A transpose is identity if permuting dimensions doesn't change which
+    // dimensions are in the same symmetric set
+    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> {
@@ -26641,7 +26719,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

…d/jax/Passes/PartialSymmetrySimplify.cpp …d/jax/Passes/PartialSymmetryAnnotate.cppsrc/enzyme_ad/jax/Passes/PartialSymmetrySimplify.cpp renamed to src/enzyme_ad/jax/Passes/PartialSymmetryAnnotate.cpp src/enzyme_ad/jax/Passes/PartialSymmetrySimplify.cpp renamed to src/enzyme_ad/jax/Passes/PartialSymmetryAnnotate.cpp

@@ -17,11 +17,11 @@
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
-#define DEBUG_TYPE "partial-symmetry-simplify"
+#define DEBUG_TYPE "partial-symmetry-annotate"
 
 namespace mlir {
 namespace enzyme {
-#define GEN_PASS_DEF_PARTIALSYMMETRYSIMPLIFYPASS
+#define GEN_PASS_DEF_PARTIALSYMMETRYANNOTATEPASS
 #include "src/enzyme_ad/jax/Passes/Passes.h.inc"
 } // namespace enzyme
 } // namespace mlir
@@ -32,9 +32,9 @@ using namespace mlir::enzyme;
 
 namespace {
 
-class PartialSymmetrySimplifyPass
-    : public enzyme::impl::PartialSymmetrySimplifyPassBase<
-          PartialSymmetrySimplifyPass> {
+class PartialSymmetryAnnotatePass
+    : public enzyme::impl::PartialSymmetryAnnotatePassBase<
+          PartialSymmetryAnnotatePass> {
 public:
   using Base::Base;
 
@@ -51,6 +51,7 @@ class PartialSymmetrySimplifyPass
 
     auto mod = getOperation();
 
+    // Annotate all operations with partial symmetry information
     mod->walk([&](Operation *op) {
       SmallVector<Attribute> partialSymmetryAttrs;
       bool anyKnown = false;
@@ -92,8 +93,6 @@ class PartialSymmetrySimplifyPass
 
       return WalkResult::advance();
     });
-
-    // TODO: do things here
   }
 };
 

src/enzyme_ad/jax/Passes/Passes.td

@@ -1077,8 +1077,8 @@ def EnzymeBatchToStableHLOPass : Pass<"enzyme-batch-to-stablehlo"> {
   ];
 }
 
-def PartialSymmetrySimplifyPass : Pass<"partial-symmetry-simplify", "ModuleOp"> {
-  let summary = "Simplify operations using partial symmetry analysis";
+def PartialSymmetryAnnotatePass : Pass<"partial-symmetry-annotate", "ModuleOp"> {
+  let summary = "Annotate operations using partial symmetry analysis";
   let dependentDialects = [
     "stablehlo::StablehloDialect",
     "enzymexla::EnzymeXLADialect",

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"];

test/lit_tests/structured_tensors/partial_symmetry.mlir

@@ -1,23 +1,25 @@
-// RUN: enzymexlamlir-opt --partial-symmetry-simplify %s | FileCheck %s
+// RUN: enzymexlamlir-opt --partial-symmetry-annotate --enzyme-hlo-generate-td="patterns=transpose_partial_symmetry_simplify" --transform-interpreter --enzyme-hlo-remove-transform %s | FileCheck %s
 
-func.func @test1() -> tensor<2x2xf32> {
+func.func @test_constant() -> tensor<2x2xf32> {
   %cst = stablehlo.constant dense<[[1.0, 2.0], [2.0, 3.0]]> : tensor<2x2xf32>
-  return %cst : tensor<2x2xf32>
+  %0 = stablehlo.transpose %cst, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  return %0 : tensor<2x2xf32>
 }
-// CHECK: func.func @test1() -> tensor<2x2xf32> {
+// CHECK: func.func @test_constant() -> tensor<2x2xf32> {
 // CHECK-NEXT: %cst = stablehlo.constant {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1]>>]} dense<{{.*}}> : tensor<2x2xf32>
 // CHECK-NEXT: return %cst : tensor<2x2xf32>
 // CHECK-NEXT: }
 
-func.func @test2() -> tensor<2x2x2x3xf32> {
+func.func @test_propagate() -> tensor<2x2x2x3xf32> {
   %cst0 = stablehlo.constant dense<[[[1.0, 2.0], [3.0, 4.0]], [[3.0, 4.0], [5.0, 6.0]]]> : tensor<2x2x2xf32>
   %cst1 = stablehlo.constant dense<[[[1.0, 2.0], [2.0, 3.0]], [[2.0, 3.0], [3.0, 4.0]]]> : tensor<2x2x2xf32>
   %0 = stablehlo.add %cst0, %cst1 : tensor<2x2x2xf32>
   %1 = stablehlo.transpose %0, dims = [0, 2, 1] : (tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
   %2 = stablehlo.broadcast_in_dim %1, dims = [1, 0, 2] : (tensor<2x2x2xf32>) -> tensor<2x2x2x3xf32>
-  return %2 : tensor<2x2x2x3xf32>
+  %3 = stablehlo.transpose %2, dims = [0, 2, 1, 3] : (tensor<2x2x2x3xf32>) -> tensor<2x2x2x3xf32>
+  return %3 : tensor<2x2x2x3xf32>
 }
-// CHECK: func.func @test2() -> tensor<2x2x2x3xf32> {
+// CHECK: func.func @test_propagate() -> tensor<2x2x2x3xf32> {
 // CHECK-NEXT: %cst = stablehlo.constant {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1]>>]} dense<{{.*}}> : tensor<2x2x2xf32>
 // CHECK-NEXT: %cst_0 = stablehlo.constant {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1, 2]>>]} dense<{{.*}}> : tensor<2x2x2xf32>
 // CHECK-NEXT: %0 = stablehlo.add %cst, %cst_0 {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1]>>]} : tensor<2x2x2xf32>
@@ -26,36 +28,36 @@ func.func @test2() -> tensor<2x2x2x3xf32> {
 // CHECK-NEXT: return %2 : tensor<2x2x2x3xf32>
 // CHECK-NEXT: }
 
-func.func @test3(%arg0: tensor<3x2x3xf32>) -> tensor<3x2x3xf32> {
+func.func @test_add_generate_symmetry(%arg0: tensor<3x2x3xf32>) -> tensor<3x2x3xf32> {
   %0 = stablehlo.transpose %arg0, dims = [2, 1, 0] : (tensor<3x2x3xf32>) -> tensor<3x2x3xf32>
   %1 = stablehlo.add %0, %arg0 : tensor<3x2x3xf32>
   return %1 : tensor<3x2x3xf32>
 }
-// CHECK: func.func @test3(%arg0: tensor<3x2x3xf32>) -> tensor<3x2x3xf32> {
+// CHECK: func.func @test_add_generate_symmetry(%arg0: tensor<3x2x3xf32>) -> tensor<3x2x3xf32> {
 // CHECK-NEXT: %0 = stablehlo.transpose %arg0, dims = [2, 1, 0] : (tensor<3x2x3xf32>) -> tensor<3x2x3xf32>
 // CHECK-NEXT: %1 = stablehlo.add %0, %arg0 {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 2]>>]} : tensor<3x2x3xf32>
 // CHECK-NEXT: return %1 : tensor<3x2x3xf32>
 // CHECK-NEXT: }
 
-func.func @test4() -> tensor<2x2xf32> {
+func.func @test_dot_propagate() -> tensor<2x2xf32> {
   %cst0 = stablehlo.constant dense<[[[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], [[2.0, 3.0, 4.0], [3.0, 4.0, 5.0]]]> : tensor<2x2x3xf32>
   %cst1 = stablehlo.constant dense<[[[1.0, 2.0], [2.0, 3.0]], [[2.0, 3.0], [3.0, 4.0]], [[2.0, 3.0], [3.0, 4.0]]]> : tensor<3x2x2xf32>
   %0 = stablehlo.dot_general %cst0, %cst1, batching_dims = [0, 1] x [1, 2], contracting_dims = [2] x [0] : (tensor<2x2x3xf32>, tensor<3x2x2xf32>) -> tensor<2x2xf32>
   return %0 : tensor<2x2xf32>
 }
-// CHECK: func.func @test4() -> tensor<2x2xf32> {
+// CHECK: func.func @test_dot_propagate() -> tensor<2x2xf32> {
 // CHECK-NEXT: %cst = stablehlo.constant {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1]>>]} dense<{{.*}}> : tensor<2x2x3xf32>
 // CHECK-NEXT: %cst_0 = stablehlo.constant {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[1, 2]>>]} dense<{{.*}}> : tensor<3x2x2xf32>
 // CHECK-NEXT: %0 = stablehlo.dot_general %cst, %cst_0, batching_dims = [0, 1] x [1, 2], contracting_dims = [2] x [0] {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[0, 1]>>]} : (tensor<2x2x3xf32>, tensor<3x2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT: return %0 : tensor<2x2xf32>
 // CHECK-NEXT: }
 
-func.func @test5(%arg0: tensor<3x3x3xf32>) -> tensor<3x3x3xf32> {
+func.func @test_dot_generate_symmetry(%arg0: tensor<3x3x3xf32>) -> tensor<3x3x3xf32> {
   %0 = stablehlo.transpose %arg0, dims = [2, 1, 0] : (tensor<3x3x3xf32>) -> tensor<3x3x3xf32>
   %1 = stablehlo.dot_general %arg0, %0, batching_dims = [1] x [1], contracting_dims = [0] x [2] : (tensor<3x3x3xf32>, tensor<3x3x3xf32>) -> tensor<3x3x3xf32>
   return %1 : tensor<3x3x3xf32>
 }
-// CHECK: func.func @test5(%arg0: tensor<3x3x3xf32>) -> tensor<3x3x3xf32> {
+// CHECK: func.func @test_dot_generate_symmetry(%arg0: tensor<3x3x3xf32>) -> tensor<3x3x3xf32> {
 // CHECK-NEXT: %0 = stablehlo.transpose %arg0, dims = [2, 1, 0] : (tensor<3x3x3xf32>) -> tensor<3x3x3xf32>
 // CHECK-NEXT: %1 = stablehlo.dot_general %arg0, %0, batching_dims = [1] x [1], contracting_dims = [0] x [2] {enzymexla.partial_symmetry = [#enzymexla.partial_symmetry<<[1, 2]>>]} : (tensor<3x3x3xf32>, tensor<3x3x3xf32>) -> tensor<3x3x3xf32>
 // CHECK-NEXT: return %1 : tensor<3x3x3xf32>