dot general case for symm (#1620)

* dot general case

* add test and fmt

* fix: update to new version

---------

Co-authored-by: Avik Pal <[email protected]>

Author: snonk

src/enzyme_ad/jax/Utils.cpp

@@ -672,8 +672,6 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
     return perm.size() == 2 && perm[0] == 1 && perm[1] == 0;
   };
 
-  // TODO: check for dot_general as well
-
   if (auto broadcastOp = dyn_cast<stablehlo::BroadcastInDimOp>(op)) {
     auto operand = broadcastOp.getOperand();
     auto operandTy = cast<RankedTensorType>(operand.getType());
@@ -710,6 +708,38 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
     }
   }
 
+  // A x (A^T) / (A^T) x A will always be symmetric
+  if (auto dotOp = dyn_cast<stablehlo::DotGeneralOp>(op)) {
+    auto dotDimNumbers = dotOp.getDotDimensionNumbers();
+    auto lhs = dotOp.getLhs();
+    auto rhs = dotOp.getRhs();
+
+    auto lhsCDims = dotDimNumbers.getLhsContractingDimensions();
+    auto rhsCDims = dotDimNumbers.getRhsContractingDimensions();
+
+    if (dotDimNumbers.getLhsBatchingDimensions().size() == 0 &&
+        dotDimNumbers.getRhsBatchingDimensions().size() == 0 &&
+        lhsCDims.size() == 1 && rhsCDims.size() == 1) {
+      if (lhs == rhs && lhsCDims[0] == rhsCDims[0]) {
+        return State::GUARANTEED;
+      }
+
+      if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
+        if (isTrueTranspose(lhsT) && lhsT.getOperand() == rhs &&
+            lhsCDims[0] == 1 - rhsCDims[0]) {
+          return State::GUARANTEED;
+        }
+      }
+
+      if (auto rhsT = rhs.getDefiningOp<stablehlo::TransposeOp>()) {
+        if (isTrueTranspose(rhsT) && rhsT.getOperand() == lhs &&
+            lhsCDims[0] == 1 - rhsCDims[0]) {
+          return State::GUARANTEED;
+        }
+      }
+    }
+  }
+
   bool recursiveCheck = false;
 
   // elementwise ops

test/lit_tests/structured_tensors/transpose_symmetric.mlir

@@ -46,6 +46,30 @@ func.func @pass3() -> tensor<3x3xf32> {
 // CHECK-NEXT:   return %2 : tensor<3x3xf32>
 // CHECK-NEXT: }
 
+func.func @pass4(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+  %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+  %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  return %2 : tensor<2x2xf32>
+}
+
+// CHECK: func.func @pass4(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+// CHECK-NEXT:   %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   return %1 : tensor<2x2xf32>
+// CHECK-NEXT: }
+
+func.func @pass5(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+  %1 = stablehlo.dot_general %arg0, %arg0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+  %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  return %2 : tensor<2x2xf32>
+}
+
+// CHECK: func.func @pass5(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+// CHECK-NEXT:   %0 = stablehlo.dot_general %arg0, %arg0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed GUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   return %0 : tensor<2x2xf32>
+// CHECK-NEXT: }
+
 func.func @fail1(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
   %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
   %1 = stablehlo.subtract %arg0, %0 : tensor<2x2xf32>
@@ -59,3 +83,17 @@ func.func @fail1(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 // CHECK-NEXT:   %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
 // CHECK-NEXT:   return %2 : tensor<2x2xf32>
 // CHECK-NEXT: }
+
+func.func @fail2(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+  %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+  %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+  return %2 : tensor<2x2xf32>
+}
+
+// CHECK: func.func @fail2(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+// CHECK-NEXT:   %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %arg0, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] {enzymexla.symmetric_matrix = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<2x2xf32>) -> tensor<2x2xf32>
+// CHECK-NEXT:   return %2 : tensor<2x2xf32>
+// CHECK-NEXT: }