add constant int/float check for symm (#1621)

* constant check

* fmt

* add splat and broadcast:

* fix: formatting

* refactor: cleanup into a single function

* add test

---------

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

Author: snonk

src/enzyme_ad/jax/Utils.cpp

@@ -592,12 +592,55 @@ SymmetricResultAnalysis initSymmetricResultAnalysis() {
   return SymmetricResultAnalysis();
 }
 
+bool checkNotEqual(APInt a, APInt b) { return a != b; }
+
+bool checkNotEqual(APFloat a, APFloat b) {
+  return a.compare(b) != llvm::APFloat::cmpEqual;
+}
+
+template <typename Ty> bool checkConstantSymmetric(DenseElementsAttr attr) {
+  if (!attr)
+    return false;
+
+  auto type = dyn_cast<RankedTensorType>(attr.getType());
+  if (!type)
+    return false;
+
+  if (type.getRank() == 0)
+    return true;
+  if (type.getRank() != 2)
+    return false;
+
+  auto shape = type.getShape();
+  int64_t rows = shape[0];
+  int64_t cols = shape[1];
+
+  if (rows != cols)
+    return false;
+  if (attr.isSplat())
+    return true;
+
+  auto values = attr.getValues<Ty>();
+  auto it = values.begin();
+
+  for (int64_t i = 0; i < rows; i++) {
+    for (int64_t j = i + 1; j < cols; j++) {
+      auto a = *(it + i * cols + j);
+      auto b = *(it + j * cols + i);
+      if (checkNotEqual(a, b))
+        return false;
+    }
+  }
+
+  return true;
+}
+
 bool SymmetricResultAnalysis::constantIntCheck(DenseElementsAttr attr) {
-  return false; // TODO
+  return checkConstantSymmetric<APInt>(attr);
 }
 
 bool SymmetricResultAnalysis::constantFloatCheck(DenseElementsAttr attr) {
-  return false; // TODO
+  return checkConstantSymmetric<APFloat>(attr);
 }
 
 SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
@@ -633,6 +676,15 @@ SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(
 
   // 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());
+    auto dims = broadcastOp.getBroadcastDimensions();
+    if (operandTy.getRank() == 0 && dims.empty()) {
+      return State::GUARANTEED;
+    }
+  }
+
   // commutative operation with A and A^T will always be symmetric
   // op(A, A^T) will also always be symmetric
   if (stablehlo::hasTraitElementwise(op) &&

test/lit_tests/structured_tensors/transpose_symmetric.mlir

@@ -26,6 +26,26 @@ func.func @pass2(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
 // CHECK-NEXT:   return %1 : tensor<2x2xf32>
 // CHECK-NEXT: }
 
+func.func @pass3() -> tensor<3x3xf32> {
+  %cst = stablehlo.constant dense<3.000000e+00> : tensor<f32>
+  %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
+  %0 = stablehlo.broadcast_in_dim %cst_0, dims = [] : (tensor<f32>) -> tensor<3x3xf32>
+  %1 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<3x3xf32>
+  %2 = stablehlo.transpose %0, dims = [1, 0] : (tensor<3x3xf32>) -> tensor<3x3xf32>
+  %3 = stablehlo.transpose %1, dims = [1, 0] : (tensor<3x3xf32>) -> tensor<3x3xf32>
+  %4 = stablehlo.add %2, %3 : tensor<3x3xf32>
+  return %4 : tensor<3x3xf32>
+}
+
+// CHECK: func.func @pass3() -> tensor<3x3xf32> {
+// CHECK-NEXT:   %cst = stablehlo.constant dense<3.000000e+00> : tensor<f32>
+// CHECK-NEXT:   %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
+// CHECK-NEXT:   %0 = stablehlo.broadcast_in_dim %cst_0, dims = [] {enzymexla.guaranteed_symmetric = true} : (tensor<f32>) -> tensor<3x3xf32>
+// CHECK-NEXT:   %1 = stablehlo.broadcast_in_dim %cst, dims = [] {enzymexla.guaranteed_symmetric = true} : (tensor<f32>) -> tensor<3x3xf32>
+// CHECK-NEXT:   %2 = stablehlo.add %0, %1 : tensor<3x3xf32>
+// CHECK-NEXT:   return %2 : tensor<3x3xf32>
+// 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>