Skip to content

Commit

Permalink
Revert "[mlir][sparse] introduce sparse_tensor.iterate operation" (#8…
Browse files Browse the repository at this point in the history 
…8953)

Reverts #88807 (merged by mistake)
  • Loading branch information
Peiming Liu authored Apr 16, 2024
1 parent 8debcf0 commit b955653
Show file tree
Hide file tree
Showing 8 changed files with 1 addition and 882 deletions.
38 changes: 0 additions & 38 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,9 @@
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/TensorEncoding.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

#include "llvm/ADT/bit.h"

//===----------------------------------------------------------------------===//
//
// Type aliases to help code be more self-documenting. Unfortunately
Expand All @@ -45,40 +41,6 @@ using Level = uint64_t;
/// including the value `ShapedType::kDynamic` (for shapes).
using Size = int64_t;

/// A simple wrapper to encode a bitset of defined (at most 64) levels.
class LevelSet {
uint64_t bits = 0;

public:
LevelSet() = default;
explicit LevelSet(uint64_t bits) : bits(bits) {}
operator uint64_t() const { return bits; }

LevelSet &set(unsigned i) {
assert(i < 64);
bits |= 1 << i;
return *this;
}

LevelSet &operator|=(LevelSet lhs) {
bits |= static_cast<uint64_t>(lhs);
return *this;
}

LevelSet &lshift(unsigned offset) {
bits = bits << offset;
return *this;
}

bool operator[](unsigned i) const {
assert(i < 64);
return (bits & (1 << i)) != 0;
}

unsigned count() const { return llvm::popcount(bits); }
bool empty() const { return bits == 0; }
};

} // namespace sparse_tensor
} // namespace mlir

Expand Down
15 changes: 0 additions & 15 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,6 @@ class SparseTensor_Attr<string name,
list<Trait> traits = []>
: AttrDef<SparseTensor_Dialect, name, traits>;

//===----------------------------------------------------------------------===//
// A simple bitset attribute wrapped over a single int64_t to encode a set of
// sparse tensor levels.
//===----------------------------------------------------------------------===//

def LevelSetAttr :
TypedAttrBase<
I64, "IntegerAttr",
And<[CPred<"::llvm::isa<::mlir::IntegerAttr>($_self)">,
CPred<"::llvm::cast<::mlir::IntegerAttr>($_self).getType().isInteger(64)">]>,
"LevelSet attribute"> {
let returnType = [{::mlir::sparse_tensor::LevelSet}];
let convertFromStorage = [{::mlir::sparse_tensor::LevelSet($_self.getValue().getZExtValue())}];
}

//===----------------------------------------------------------------------===//
// These attributes are just like `IndexAttr` except that they clarify whether
// the index refers to a dimension (an axis of the semantic tensor) or a level
Expand Down
152 changes: 1 addition & 151 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@ include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.td"
include "mlir/Dialect/SparseTensor/IR/SparseTensorInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/LoopLikeInterface.td"

//===----------------------------------------------------------------------===//
// Base class.
Expand Down Expand Up @@ -1279,7 +1277,7 @@ def SparseTensor_SelectOp : SparseTensor_Op<"select", [Pure, SameOperandsAndResu

def SparseTensor_YieldOp : SparseTensor_Op<"yield", [Pure, Terminator,
ParentOneOf<["BinaryOp", "UnaryOp", "ReduceOp", "SelectOp",
"ForeachOp", "IterateOp"]>]>,
"ForeachOp"]>]>,
Arguments<(ins Variadic<AnyType>:$results)> {
let summary = "Yield from sparse_tensor set-like operations";
let description = [{
Expand Down Expand Up @@ -1432,154 +1430,6 @@ def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// Sparse Tensor Iteration Operations.
//===----------------------------------------------------------------------===//

def ExtractIterSpaceOp : SparseTensor_Op<"extract_iteration_space",
[Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {

let arguments = (ins AnySparseTensor:$tensor,
Optional<AnySparseIterator>:$parentIter,
LevelAttr:$loLvl, LevelAttr:$hiLvl);

let results = (outs AnySparseIterSpace:$resultSpace);

let summary = "Extract an iteration space from a sparse tensor between certain levels";
let description = [{
Extracts a `!sparse_tensor.iter_space` from a sparse tensor between
certian (consecutive) levels.

`tensor`: the input sparse tensor that defines the iteration space.
`parentIter`: the iterator for the previous level, at which the iteration space
at the current levels will be extracted.
`loLvl`, `hiLvl`: the level range between [loLvl, hiLvl) in the input tensor that
the returned iteration space covers. `hiLvl - loLvl` defines the dimension of the
iteration space.

Example:
```mlir
// Extracts a 1-D iteration space from a COO tensor at level 1.
%space = sparse_tensor.iteration.extract_space %sp at %it1 lvls = 1
: tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0>
```
}];


let extraClassDeclaration = [{
std::pair<Level, Level> getLvlRange() {
return std::make_pair(getLoLvl(), getHiLvl());
}
unsigned getSpaceDim() {
return getHiLvl() - getLoLvl();
}
ArrayRef<::mlir::sparse_tensor::LevelType> getSpaceLvlTypes() {
return getResultSpace().getType().getLvlTypes();
}
}];

let hasVerifier = 1;
let assemblyFormat = "$tensor (`at` $parentIter^)? `lvls` `=` custom<LevelRange>($loLvl, $hiLvl) "
" attr-dict `:` type($tensor) (`,` type($parentIter)^)?";
}

def IterateOp : SparseTensor_Op<"iterate",
[RecursiveMemoryEffects, RecursivelySpeculatable,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInitsMutable", "getLoopResults", "getRegionIterArgs",
"getYieldedValuesMutable"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>,
SingleBlockImplicitTerminator<"sparse_tensor::YieldOp">]> {

let arguments = (ins AnySparseIterSpace:$iterSpace,
Variadic<AnyType>:$initArgs,
LevelSetAttr:$crdUsedLvls);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);

let summary = "Iterate over a sparse iteration space";
let description = [{
The `sparse_tensor.iterate` operations represents a loop over the
provided iteration space extracted from a specific sparse tensor.
The operation defines an SSA value for a sparse iterator that points
to the current stored element in the sparse tensor and SSA values
for coordinates of the stored element. The coordinates are always
converted to `index` type despite of the underlying sparse tensor
storage. When coordinates are not used, the SSA values can be skipped
by `_` symbols, which usually leads to simpler generated code after
sparsification. For example:

```mlir
// The coordinate for level 0 is not used when iterating over a 2-D
// iteration space.
%sparse_tensor.iterate %iterator in %space at(_, %crd_1)
: !sparse_tensor.iter_space<#CSR, lvls = 0 to 2>
```

`sparse_tensor.iterate` can also operate on loop-carried variables
and returns the final values after loop termination.
The initial values of the variables are passed as additional SSA operands
to the iterator SSA value and used coordinate SSA values mentioned
above. The operation region has an argument for the iterator, variadic
arguments for specified (used) coordiates and followed by one argument
for each loop-carried variable, representing the value of the variable
at the current iteration.
The body region must contain exactly one block that terminates with
`sparse_tensor.yield`.

`sparse_tensor.iterate` results hold the final values after the last
iteration. If the `sparse_tensor.iterate` defines any values, a yield
must be explicitly present.
The number and types of the `sparse_tensor.iterate` results must match
the initial values in the iter_args binding and the yield operands.


A nested `sparse_tensor.iterate` example that prints all the coordinates
stored in the sparse input:

```mlir
func.func @nested_iterate(%sp : tensor<4x8xf32, #COO>) {
// Iterates over the first level of %sp
%l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 : tensor<4x8xf32, #COO>
%r1 = sparse_tensor.iterate %it1 in %l1 at (%crd0)
: !sparse_tensor.iter_space<#COO, lvls = 0 to 1> {
// Iterates over the second level of %sp
%l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1
: tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1>
%r2 = sparse_tensor.iterate %it2 in %l2 at (crd1)
: !sparse_tensor.iter_space<#COO, lvls = 1 to 2> {
vector.print %crd0 : index
vector.print %crd1 : index
}
}
}

```
}];

let extraClassDeclaration = [{
unsigned getSpaceDim() {
return getIterSpace().getType().getSpaceDim();
}
BlockArgument getIterator() {
return getRegion().getArguments().front();
}
Block::BlockArgListType getCrds() {
// The first block argument is iterator, the remaining arguments are
// referenced coordinates.
return getRegion().getArguments().slice(1, getCrdUsedLvls().count());
}
unsigned getNumRegionIterArgs() {
return getRegion().getArguments().size() - 1 - getCrdUsedLvls().count();
}
}];

let hasVerifier = 1;
let hasRegionVerifier = 1;
let hasCustomAssemblyFormat = 1;
}

//===----------------------------------------------------------------------===//
// Sparse Tensor Debugging and Test-Only Operations.
//===----------------------------------------------------------------------===//
Expand Down
95 changes: 0 additions & 95 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,99 +72,4 @@ def SparseTensorStorageSpecifier
: Type<CPred<"::llvm::isa<::mlir::sparse_tensor::StorageSpecifierType>($_self)">, "metadata",
"::mlir::sparse_tensor::StorageSpecifierType">;

//===----------------------------------------------------------------------===//
// Sparse Tensor Iteration Types.
//===----------------------------------------------------------------------===//

def SparseTensor_IterSpace : SparseTensor_Type<"IterSpace"> {
let mnemonic = "iter_space";

let description = [{
A sparse iteration space that represents an abstract N-D (sparse) iteration space
extracted from a sparse tensor.

Examples:

```mlir
// An iteration space extracted from a CSR tensor between levels [0, 2).
!iter_space<#CSR, lvls = 0 to 2>
```
}];

let parameters = (ins
SparseTensorEncodingAttr : $encoding,
"Level" : $loLvl,
"Level" : $hiLvl
);

let extraClassDeclaration = [{
/// The the dimension of the iteration space.
unsigned getSpaceDim() const {
return getHiLvl() - getLoLvl();
}

/// Get the level types for the iteration space.
ArrayRef<LevelType> getLvlTypes() const {
return getEncoding().getLvlTypes().slice(getLoLvl(), getSpaceDim());
}

/// Whether the iteration space is unique (i.e., no duplicated coordinate).
bool isUnique() {
return !getLvlTypes().back().isa<LevelPropNonDefault::Nonunique>();
}

/// Get the corresponding iterator type.
::mlir::sparse_tensor::IteratorType getIteratorType() const;
}];

let assemblyFormat="`<` $encoding `,` `lvls` `=` custom<LevelRange>($loLvl, $hiLvl) `>`";
}

def SparseTensor_Iterator : SparseTensor_Type<"Iterator"> {
let mnemonic = "iterator";

let description = [{
An iterator that points to the current element in the corresponding iteration space.

Examples:

```mlir
// An iterator that iterates over a iteration space of type `!iter_space<#CSR, lvls = 0 to 2>`
!iterator<#CSR, lvls = 0 to 2>
```
}];

let parameters = (ins
SparseTensorEncodingAttr : $encoding,
"Level" : $loLvl,
"Level" : $hiLvl
);

let extraClassDeclaration = [{
/// Get the corresponding iteration space type.
::mlir::sparse_tensor::IterSpaceType getIterSpaceType() const;

unsigned getSpaceDim() const { return getIterSpaceType().getSpaceDim(); }
ArrayRef<LevelType> getLvlTypes() const { return getIterSpaceType().getLvlTypes(); }
bool isUnique() { return getIterSpaceType().isUnique(); }
}];

let assemblyFormat="`<` $encoding `,` `lvls` `=` custom<LevelRange>($loLvl, $hiLvl) `>`";
}

def IsSparseSparseIterSpaceTypePred
: CPred<"::llvm::isa<::mlir::sparse_tensor::IterSpaceType>($_self)">;

def IsSparseSparseIteratorTypePred
: CPred<"::llvm::isa<::mlir::sparse_tensor::IteratorType>($_self)">;

def AnySparseIterSpace
: Type<IsSparseSparseIterSpaceTypePred, "sparse iteration space",
"::mlir::sparse_tensor::IterSpaceType">;

def AnySparseIterator
: Type<IsSparseSparseIteratorTypePred, "sparse iterator",
"::mlir::sparse_tensor::IteratorType">;


#endif // SPARSETENSOR_TYPES
Loading

0 comments on commit b955653

Please sign in to comment.