feat: more propagation rules

Author: avik-pal

src/enzyme_ad/jax/Analysis/StructuredMatrixAnalysis.cpp

@@ -33,8 +33,6 @@ StructuredSparsityPattern::StructuredSparsityPattern(Value v) {
     return;
   }
 
-  llvm::errs() << "TODO: structured sparsity pattern not implemented for " << v
-               << "\n";
   setUnknown();
   return;
 }
@@ -130,9 +128,8 @@ StructuredSparsityPattern::meet(const StructuredSparsityPattern &lhs,
   if (rhs.kind == StructuredSparsityKind::Unknown)
     return lhs;
 
-  // for all other cases, we take the min of the bandwidths and refine
-  auto lb = std::min(lhs.lowerBandwidth, rhs.lowerBandwidth);
-  auto ub = std::min(lhs.upperBandwidth, rhs.upperBandwidth);
+  auto lb = std::max(lhs.lowerBandwidth, rhs.lowerBandwidth);
+  auto ub = std::max(lhs.upperBandwidth, rhs.upperBandwidth);
   auto newPattern = StructuredSparsityPattern(lb, ub);
   newPattern.refineKind();
   return newPattern;
@@ -151,13 +148,21 @@ StructuredSparsityPattern::join(const StructuredSparsityPattern &lhs,
       rhs.kind == StructuredSparsityKind::Unknown)
     return StructuredSparsityPattern(StructuredSparsityKind::Unknown);
 
-  auto lb = std::max(lhs.lowerBandwidth, rhs.lowerBandwidth);
-  auto ub = std::max(lhs.upperBandwidth, rhs.upperBandwidth);
+  auto lb = std::min(lhs.lowerBandwidth, rhs.lowerBandwidth);
+  auto ub = std::min(lhs.upperBandwidth, rhs.upperBandwidth);
   auto newPattern = StructuredSparsityPattern(lb, ub);
   newPattern.refineKind();
   return newPattern;
 }
 
+StructuredSparsityPattern StructuredSparsityPattern::propagateTranspose(
+    const StructuredSparsityPattern &op) {
+  auto newPattern = StructuredSparsityPattern(op.upperBandwidth,
+                                              op.lowerBandwidth);
+  newPattern.refineKind();
+  return newPattern;
+}
+
 void StructuredSparsityPattern::print(raw_ostream &os) const {
   switch (kind) {
   case StructuredSparsityKind::Unknown:
@@ -413,6 +418,31 @@ void StructuredMatrixType::print(raw_ostream &os) const {
   os << ")";
 }
 
+StructuredMatrixType StructuredMatrixType::propagateTranspose(
+    const StructuredMatrixType &op) {
+  return StructuredMatrixType(
+      StructuredSparsityPattern::propagateTranspose(op.sparsityPattern),
+      op.valueProperties);
+}
+
+StructuredMatrixType StructuredMatrixType::propagateAdd(
+    const StructuredMatrixType &lhs, const StructuredMatrixType &rhs) {
+  ValueProperties valProps;
+  // TODO: If one is unit diag and other is zeros, we can propagate the other
+  // to the unit diag
+  if (lhs.getProperties().isSymmetric() && rhs.getProperties().isSymmetric()) {
+    valProps.set(ValueProperty::Symmetric);
+  }
+  if (lhs.getProperties().isBroadcastedScalar() &&
+      rhs.getProperties().isBroadcastedScalar()) {
+    valProps.set(ValueProperty::BroadcastedScalar);
+  }
+
+  return StructuredMatrixType(
+      StructuredSparsityPattern::meet(lhs.sparsityPattern, rhs.sparsityPattern),
+      valProps);
+}
+
 //===----------------------------------------------------------------------===//
 // Lattice Element
 //===----------------------------------------------------------------------===//
@@ -465,13 +495,46 @@ void StructuredMatrixAnalysis::setToEntryState(
 LogicalResult StructuredMatrixAnalysis::visitOperation(
     Operation *op, ArrayRef<const StructuredMatrixLattice *> operands,
     ArrayRef<StructuredMatrixLattice *> results) {
+  SmallVector<bool> updatedProps(results.size(), false);
+  SmallVector<StructuredMatrixType> propagatedProps(results.size());
+
+  // transpose
+  if (auto transposeOp = dyn_cast<stablehlo::TransposeOp>(op)) {
+    updatedProps[0] = true;
+    propagatedProps[0] = StructuredMatrixType::propagateTranspose(
+        operands[0]->getValue());
+  }
+
+  // elementwise
+  /// add
+  if (auto addOp = dyn_cast<stablehlo::AddOp>(op)) {
+    updatedProps[0] = true;
+    propagatedProps[0] = StructuredMatrixType::propagateAdd(
+        operands[0]->getValue(), operands[1]->getValue());
+  }
+
+  /// mul
+
+  // finalize
+  for (size_t i = 0; i < results.size(); i++) {
+    if (updatedProps[i]) {
+      results[i]->setValue(
+          StructuredMatrixType::join(results[i]->getValue(), propagatedProps[i]));
+    }
+  }
+
 
   llvm::errs() << "Visiting operation " << *op << "\n";
   for (auto operand : operands) {
     llvm::errs() << "    operand: ";
     operand->getValue().print(llvm::errs());
     llvm::errs() << "\n";
   }
+  for (auto result : results) {
+    llvm::errs() << "    result: ";
+    result->getValue().print(llvm::errs());
+    llvm::errs() << "\n";
+  }
   llvm::errs() << "\n";
 
   return success();

src/enzyme_ad/jax/Analysis/StructuredMatrixAnalysis.h

@@ -88,6 +88,10 @@ class StructuredSparsityPattern {
     return os;
   }
 
+  // propagation rules
+  static StructuredSparsityPattern propagateTranspose(
+      const StructuredSparsityPattern &op);
+
 private:
   void initializeBandwidths();
   void refineKind();
@@ -203,7 +207,14 @@ class StructuredMatrixType {
     return os;
   }
 
-  // TODO: propagation rules probably goes in here
+  // propagation rules
+  static StructuredMatrixType propagateTranspose(const StructuredMatrixType &op);
+
+  static StructuredMatrixType propagateAdd(const StructuredMatrixType &lhs,
+                                           const StructuredMatrixType &rhs);
+
+  static StructuredMatrixType propagateMultiply(const StructuredMatrixType &lhs,
+                                                const StructuredMatrixType &rhs);
 
   // TODO: implement queries that check both the sparsity pattern and value
   // properties and return specific matrix kinds