[Mlir-commits] [mlir] 5412ed4 - [mlir][sparse] Factoring out Merger::expIsTensor

[wren romano]

Author: wren romano
Date: 2023-03-08T16:42:34-08:00
New Revision: 5412ed44b5887563a17203d5ed5796dbffce3577

URL: https://github.com/llvm/llvm-project/commit/5412ed44b5887563a17203d5ed5796dbffce3577
DIFF: https://github.com/llvm/llvm-project/commit/5412ed44b5887563a17203d5ed5796dbffce3577.diff

LOG: [mlir][sparse] Factoring out Merger::expIsTensor

Depends On D145611

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145620

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
    mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
index c96f76ee1c094..4ad069d18fd81 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
@@ -267,6 +267,11 @@ class Merger {
     return ldx >= numNativeLoops;
   }
 
+  /// Returns true if the expression is `(kTensor t)`.
+  bool expIsTensor(unsigned e, unsigned t) const {
+    return tensorExps[e].kind == kTensor && tensorExps[e].tensor == t;
+  }
+
   /// Returns true if the expression contains the `t` as an operand.
   bool expContainsTensor(unsigned e, unsigned t) const;
 

diff  --git a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp b/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
index 4614a103b77e2..0fd75476db25f 100644
--- a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
@@ -318,7 +318,7 @@ unsigned Merger::optimizeSet(unsigned s0) {
     if (p0 != p1) {
       // Is this a straightforward copy?
       unsigned e = latPoints[p1].exp;
-      if (tensorExps[e].kind == kTensor && tensorExps[e].tensor == outTensor)
+      if (expIsTensor(e, outTensor))
         continue;
       // Conjunction already covered?
       for (unsigned p2 : latSets[s]) {
@@ -405,15 +405,14 @@ bool Merger::expContainsTensor(unsigned e, unsigned t) const {
     return false;
   case ExpArity::kUnary: {
     unsigned op = tensorExps[e].children.e0;
-    if (tensorExps[op].kind == kTensor && tensorExps[op].tensor == t)
+    if (expIsTensor(op, t))
       return true;
     return expContainsTensor(op, t);
   }
   case ExpArity::kBinary: {
     unsigned op1 = tensorExps[e].children.e0;
     unsigned op2 = tensorExps[e].children.e1;
-    if ((tensorExps[op1].kind == kTensor && tensorExps[op1].tensor == t) ||
-        (tensorExps[op2].kind == kTensor && tensorExps[op2].tensor == t))
+    if (expIsTensor(op1, t) || expIsTensor(op2, t))
       return true;
     return expContainsTensor(op1, t) || expContainsTensor(op2, t);
   }