Skip to content

[mlir][sparse] remove COO test from trait and encoding #73733

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 29, 2023
Merged

[mlir][sparse] remove COO test from trait and encoding #73733

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d2f6ab3
[mlir][sparse] remove COO test from trait and encoding
aartbik Nov 29, 2023
b9183fb
DCE
aartbik Nov 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 0 additions & 15 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -467,33 +464,21 @@ 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()">;

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]>;

//===----------------------------------------------------------------------===//
Expand Down
15 changes: 7 additions & 8 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down
12 changes: 6 additions & 6 deletions mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -1664,14 +1660,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();
}

Expand Down