feat: cleaner printing

Author: avik-pal

src/enzyme_ad/jax/Analysis/StructuredMatrixAnalysis.h

@@ -67,6 +67,7 @@ class StructuredSparsityPattern {
     refineKind();
   }
 
+  StructuredSparsityKind getKind() const { return kind; }
   int64_t getLowerBandwidth() const { return lowerBandwidth; }
   int64_t getUpperBandwidth() const { return upperBandwidth; }
 

src/enzyme_ad/jax/Passes/StructuredMatrixSimplify.cpp

@@ -70,31 +70,71 @@ class StructuredMatrixSimplifyPass
           continue;
         }
 
-        anyKnown = true;
+        if (state->getValue().getSparsityPattern().getKind() !=
+            mlir::structure_analysis::StructuredSparsityKind::Unknown) {
+          anyKnown = true;
+        }
+
+        enzymexla::StructuredSparsityKind ssKind;
+        switch (state->getValue().getSparsityPattern().getKind()) {
+        case mlir::structure_analysis::StructuredSparsityKind::Unknown:
+          ssKind = enzymexla::StructuredSparsityKind::Unknown;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::Dense:
+          ssKind = enzymexla::StructuredSparsityKind::Dense;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::Band:
+          ssKind = enzymexla::StructuredSparsityKind::Band;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::UpperTriangular:
+          ssKind = enzymexla::StructuredSparsityKind::UpperTriangular;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::UpperBidiagonal:
+          ssKind = enzymexla::StructuredSparsityKind::UpperBidiagonal;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::LowerTriangular:
+          ssKind = enzymexla::StructuredSparsityKind::LowerTriangular;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::LowerBidiagonal:
+          ssKind = enzymexla::StructuredSparsityKind::LowerBidiagonal;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::Tridiagonal:
+          ssKind = enzymexla::StructuredSparsityKind::Tridiagonal;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::Diagonal:
+          ssKind = enzymexla::StructuredSparsityKind::Diagonal;
+          break;
+        case mlir::structure_analysis::StructuredSparsityKind::Empty:
+          ssKind = enzymexla::StructuredSparsityKind::Empty;
+          break;
+        }
 
-        // TODO: get structured sparsity kind
         auto structuredSparsityKind =
             enzymexla::StructuredSparsityPatternAttr::get(
-                mod.getContext(), enzymexla::StructuredSparsityKind::Unknown,
+                mod.getContext(), ssKind,
                 state->getValue().getSparsityPattern().getLowerBandwidth(),
                 state->getValue().getSparsityPattern().getUpperBandwidth());
 
         SmallVector<enzymexla::StructuredValueProperty>
             structuredValueProperties;
         auto valueProperties = state->getValue().getProperties();
         if (valueProperties.hasUnitDiagonal()) {
+          anyKnown = true;
           structuredValueProperties.push_back(
               enzymexla::StructuredValueProperty::UnitDiagonal);
         }
         if (valueProperties.isSymmetric()) {
+          anyKnown = true;
           structuredValueProperties.push_back(
               enzymexla::StructuredValueProperty::Symmetric);
         }
         if (valueProperties.isHermitian()) {
+          anyKnown = true;
           structuredValueProperties.push_back(
               enzymexla::StructuredValueProperty::Hermitian);
         }
         if (valueProperties.isBroadcastedScalar()) {
+          anyKnown = true;
           structuredValueProperties.push_back(
               enzymexla::StructuredValueProperty::BroadcastedScalar);
         }