feat: broadcast check

Author: avik-pal

src/enzyme_ad/jax/Analysis/StructuredMatrixAnalysis.cpp

@@ -13,6 +13,8 @@
 #include "mlir/Support/LLVM.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "stablehlo/dialect/StablehloOps.h"
+
 using namespace mlir;
 using namespace mlir::dataflow;
 
@@ -75,9 +77,6 @@ void StructuredSparsityPattern::initializeBandwidths() {
 }
 
 void StructuredSparsityPattern::refineKind() {
-  if (kind != StructuredSparsityKind::Band)
-    return;
-
   if (lowerBandwidth == 0) {
     if (upperBandwidth == 0) {
       kind = StructuredSparsityKind::Diagonal;
@@ -215,15 +214,26 @@ ValueProperties::ValueProperties(Value v) {
   if (matchPattern(v, m_Constant(&denseAttr))) {
     auto props = getPropertiesFromDenseAttr(denseAttr);
     setFlags(props.getFlags());
-    llvm::errs() << "v: " << v << " properties: ";
-    this->print(llvm::errs());
-    llvm::errs() << "\n";
     return;
   }
 
-  // TODO: symmetric checks Utils.cpp:688
+  auto defOp = v.getDefiningOp();
+  if (!defOp)
+    return;
+
+  // comm_op(A, A^T) will always be symmetric
 
-  // TODO: broadcasted scalar
+  // A x A^T will always be symmetric
+
+  if (auto bcastOp = dyn_cast<stablehlo::BroadcastInDimOp>(defOp)) {
+    auto operand = bcastOp.getOperand();
+    if (cast<RankedTensorType>(operand.getType()).getRank() == 0) { // bcast(scalar)
+      if (matchPattern(operand, m_One())) // bcast(1)
+        set(ValueProperty::UnitDiagonal);
+      set(ValueProperty::BroadcastedScalar);
+      set(ValueProperty::Symmetric);
+    }
+  }
 
   // TODO: unit diagonal
   //       - iota scatter with constant
@@ -444,9 +454,5 @@ LogicalResult StructuredMatrixAnalysis::visitOperation(
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// Structure Originators
-//===----------------------------------------------------------------------===//
-
 } // namespace structure_analysis
 } // namespace mlir

src/enzyme_ad/jax/Analysis/StructuredMatrixAnalysis.h

@@ -254,9 +254,5 @@ class StructuredMatrixAnalysis
                  ArrayRef<StructuredMatrixLattice *> results) override;
 };
 
-//===----------------------------------------------------------------------===//
-// Structure Originators
-//===----------------------------------------------------------------------===//
-
 } // namespace structure_analysis
 } // namespace mlir