add test and fmt

Author: snonk

src/enzyme_ad/jax/Utils.cpp

@@ -710,14 +710,14 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
     }
   }
 
-    // A(A^T) will always be symmetric
+  // A(A^T) will always be symmetric
   if (auto dotOp = dyn_cast_or_null<stablehlo::DotGeneralOp>(op)) {
     auto lhs = dotOp.getLhs();
     auto rhs = dotOp.getRhs();
 
     auto lhsDims = dotOp.getDotDimensionNumbers().getLhsContractingDimensions();
     auto rhsDims = dotOp.getDotDimensionNumbers().getRhsContractingDimensions();
-    
+
     if (auto lhsType = dyn_cast_or_null<ShapedType>(lhs.getType());
         lhsType && lhsType.hasRank() && lhsType.getRank() == 2) {
 
@@ -733,8 +733,8 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
           if (rhsInput == lhs && isTrueTranspose(rhsT))
             return State::GUARANTEED;
         }
-    
-        if (auto lhsT = dyn_cast_or_null<stablehlo::TransposeOp>(lhs.getDefiningOp())) {
+
+        if (auto lhsT = lhs.getDefiningOp<stablehlo::TransposeOp>()) {
           auto lhsInput = lhsT.getOperand();
           if (lhsInput == rhs && isTrueTranspose(lhsT))
             return State::GUARANTEED;

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.guaranteed_symmetric = true} : (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.guaranteed_symmetric = true} : (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.guaranteed_symmetric = false} : (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: }