[Mlir-commits] [mlir] 250f689 - [mlir][sparse] convenience util to test for "sparse" op

[Aart Bik]

Author: Aart Bik
Date: 2023-03-02T13:08:50-08:00
New Revision: 250f6899c3d52a797268adde24083c439b508b2e

URL: https://github.com/llvm/llvm-project/commit/250f6899c3d52a797268adde24083c439b508b2e
DIFF: https://github.com/llvm/llvm-project/commit/250f6899c3d52a797268adde24083c439b508b2e.diff

LOG: [mlir][sparse] convenience util to test for "sparse" op

Reviewed By: cota

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
index 8718697e0cc85..d61aaed3b33fb 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
@@ -116,6 +116,17 @@ RankedTensorType getCOOFromTypeWithOrdering(RankedTensorType src,
 
 RankedTensorType getCOOFromType(RankedTensorType src, bool ordered);
 
+/// Returns true iff MLIR operand has any sparse operand or result.
+inline bool hasAnySparseOperandOrResult(Operation *op) {
+  bool anySparseIn = llvm::any_of(op->getOperands().getTypes(), [](Type t) {
+    return getSparseTensorEncoding(t) != nullptr;
+  });
+  bool anySparseOut = llvm::any_of(op->getResults().getTypes(), [](Type t) {
+    return getSparseTensorEncoding(t) != nullptr;
+  });
+  return anySparseIn || anySparseOut;
+}
+
 //
 // Reordering.
 //