[Mlir-commits] [mlir] [mlir][sparse] remove COO test from trait and encoding (PR #73733)
Aart Bik
llvmlistbot at llvm.org
Tue Nov 28 17:33:14 PST 2023
https://github.com/aartbik updated https://github.com/llvm/llvm-project/pull/73733
>From d2f6ab38158cd31445766ba9b357a1338e527988 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Tue, 28 Nov 2023 17:27:44 -0800
Subject: [PATCH 1/2] [mlir][sparse] remove COO test from trait and encoding
This is a minor step towards moving ALL COO related
tests into the SparseTensorType class rather than
having it all over the place (with risk of becoming
inconsistent). Next revision will move ALL COO related
methods into this class.
---
.../SparseTensor/IR/SparseTensorAttrDefs.td | 12 ------------
.../Dialect/SparseTensor/IR/SparseTensorOps.td | 15 +++++++--------
.../SparseTensor/IR/SparseTensorDialect.cpp | 8 ++++++--
3 files changed, 13 insertions(+), 22 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 1e62d9935d63c32..13f97994a57b8e0 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -467,10 +467,6 @@ def SparseTensorStorageSpecifierKindAttr
def IsSparseTensorPred
: CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
-def IsCOOPred
- : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self) && "
- " ::mlir::sparse_tensor::getSparseTensorEncoding($_self).isCOO()">;
-
def IsSparseTensorSlicePred
: CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self) && "
" ::mlir::sparse_tensor::getSparseTensorEncoding($_self).isSlice()">;
@@ -478,22 +474,14 @@ def IsSparseTensorSlicePred
class SparseTensorOf<list<Type> allowedTypes>
: TensorOf<allowedTypes, [IsSparseTensorPred], "sparse tensor">;
-class COOSparseTensorOf<list<Type> allowedTypes>
- : TensorOf<allowedTypes, [IsCOOPred], "COO sparse tensor">;
-
class SparseTensorSliceOf<list<Type> allowedTypes>
: TensorOf<allowedTypes, [IsSparseTensorSlicePred], "sparse tensor slice">;
-class RankedSparseTensorOf<list<Type> allowedTypes>
- : RankedTensorOf<allowedTypes, [IsSparseTensorPred], "ranked sparse tensor">;
-
class ScalarLikeOf<list<Type> allowedTypes>
: AnyTypeOf<[0DTensorOf<allowedTypes>, AnyTypeOf<allowedTypes>], "scalar like">;
def AnySparseTensor : SparseTensorOf<[AnyType]>;
-def AnyCOOSparseTensor : COOSparseTensorOf<[AnyType]>;
def AnySparseTensorSlice : SparseTensorSliceOf<[AnyType]>;
-def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
def AnyIndexingScalarLike : ScalarLikeOf<[AnySignlessIntegerOrIndex]>;
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index 78031f28949a9e5..3127cf1b1bcf696 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -921,10 +921,9 @@ def SparseTensor_SortOp : SparseTensor_Op<"sort">,
let summary = "Sorts the arrays in xs and ys lexicographically on the "
"integral values found in the xs list";
let description = [{
- Sparse_tensor.sort sort the `xs` values along with some `ys` values
- that are put in a single linear buffer `xy`.
- The affine map attribute `perm_map` specifies the permutation to be applied on
- the `xs` before comparison, the rank of the permutation map
+ Sorts the `xs` values along with some `ys` values that are put in a single linear
+ buffer `xy`. The affine map attribute `perm_map` specifies the permutation to be
+ applied on the `xs` before comparison, the rank of the permutation map
also specifies the number of `xs` values in `xy`.
The optional index attribute `ny` provides the number of `ys` values in `xy`.
When `ny` is not explicitly specified, its value is 0.
@@ -950,14 +949,14 @@ def SparseTensor_SortOp : SparseTensor_Op<"sort">,
}
def SparseTensor_ReorderCOOOp : SparseTensor_Op<"reorder_coo", [Pure]>,
- Arguments<(ins AnyCOOSparseTensor: $input_coo,
+ Arguments<(ins AnySparseTensor: $input_coo,
SparseTensorSortKindAttr:$algorithm)>,
- Results<(outs AnyCOOSparseTensor: $result_coo)> {
+ Results<(outs AnySparseTensor: $result_coo)> {
let summary = "Reorder the input COO such that it has the the same order as "
"the output COO";
let description = [{
- sparse_tensor.reorder_coo reorder input COO to the same order as specified by
- the output format. E.g., reorder an unordered COO into an ordered one.
+ Reorders the input COO to the same order as specified by the output format.
+ E.g., reorder an unordered COO into an ordered one.
The input and result COO tensor must have the same element type, position type and
coordinate type. At the moment, the operation also only supports ordering
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 74d2fd5fd9f829c..51e4b45525a410b 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1664,14 +1664,18 @@ LogicalResult ReorderCOOOp::verify() {
SparseTensorType srcStt = getSparseTensorType(getInputCoo());
SparseTensorType dstStt = getSparseTensorType(getResultCoo());
+ if (!isCOOType(srcStt.getEncoding(), 0, /*isUnique=*/true) ||
+ !isCOOType(dstStt.getEncoding(), 0, /*isUnique=*/true))
+ emitError("Unexpected non-COO sparse tensors");
+
if (!srcStt.hasSameDimToLvl(dstStt))
emitError("Unmatched dim2lvl map between input and result COO");
if (srcStt.getPosType() != dstStt.getPosType() ||
srcStt.getCrdType() != dstStt.getCrdType() ||
- srcStt.getElementType() != dstStt.getElementType()) {
+ srcStt.getElementType() != dstStt.getElementType())
emitError("Unmatched storage format between input and result COO");
- }
+
return success();
}
>From b9183fb1e5816f42eae765744177fd4e32d6ddf4 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Tue, 28 Nov 2023 17:32:26 -0800
Subject: [PATCH 2/2] DCE
---
.../mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td | 3 ---
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp | 4 ----
2 files changed, 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 13f97994a57b8e0..12c1068ae1f546e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -382,9 +382,6 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
/// the null encoding (since dense-tensors are always all-dense).
bool isAllDense() const;
- /// Returns true if it is a sparse tensor encoding in COO format.
- bool isCOO() const;
-
/// Returns true if every level is ordered. Also returns true for
/// the null encoding (since dense-tensors are always all-ordered).
bool isAllOrdered() const;
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 51e4b45525a410b..d4f8afdd62f2383 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -316,10 +316,6 @@ bool SparseTensorEncodingAttr::isAllDense() const {
return !getImpl() || llvm::all_of(getLvlTypes(), isDenseLT);
}
-bool SparseTensorEncodingAttr::isCOO() const {
- return getImpl() && isCOOType(*this, 0, true);
-}
-
bool SparseTensorEncodingAttr::isAllOrdered() const {
return !getImpl() || llvm::all_of(getLvlTypes(), isOrderedLT);
}
More information about the Mlir-commits
mailing list