move COOSegments

Author: yinying-lisa-li

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

@@ -41,6 +41,19 @@ using Level = uint64_t;
 /// including the value `ShapedType::kDynamic` (for shapes).
 using Size = int64_t;
 
+/// A simple structure that encodes a range of levels in the sparse tensors
+/// that forms a COO segment.
+struct COOSegment {
+  std::pair<Level, Level> lvlRange; // [low, high)
+  bool isSoA;
+
+  bool isAoS() const { return !isSoA; }
+  bool isSegmentStart(Level l) const { return l == lvlRange.first; }
+  bool inSegment(Level l) const {
+    return l >= lvlRange.first && l < lvlRange.second;
+  }
+};
+
 } // namespace sparse_tensor
 } // namespace mlir
 

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

@@ -507,22 +507,9 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     ValueRange translateCrds(::mlir::OpBuilder &builder, ::mlir::Location loc, ::mlir::ValueRange crds, ::mlir::sparse_tensor::CrdTransDirectionKind) const;
 
     //
-    // COO struct and methods.
+    // COO methods.
     //
 
-    /// A simple structure that encodes a range of levels in the sparse tensors
-    /// that forms a COO segment.
-    struct COOSegment {
-      std::pair<Level, Level> lvlRange; // [low, high)
-      bool isSoA;
-
-      bool isAoS() const { return !isSoA; }
-      bool isSegmentStart(Level l) const { return l == lvlRange.first; }
-      bool inSegment(Level l) const {
-        return l >= lvlRange.first && l < lvlRange.second;
-      }
-    };
-
     /// Returns the starting level of this sparse tensor type for a
     /// trailing COO region that spans **at least** two levels. If
     /// no such COO region is found, then returns the level-rank.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h

@@ -356,7 +356,7 @@ class SparseTensorType {
   RankedTensorType getCOOType(bool ordered) const;
 
   /// Returns a list of COO segments in the sparse tensor types.
-  SmallVector<SparseTensorEncodingAttr::COOSegment> getCOOSegments() const {
+  SmallVector<COOSegment> getCOOSegments() const {
     return getEncoding().getCOOSegments();
   }
 

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

@@ -104,8 +104,7 @@ void StorageLayout::foreachField(
         callback) const {
   const auto lvlTypes = enc.getLvlTypes();
   const Level lvlRank = enc.getLvlRank();
-  SmallVector<SparseTensorEncodingAttr::COOSegment> cooSegs =
-      enc.getCOOSegments();
+  SmallVector<COOSegment> cooSegs = enc.getCOOSegments();
   FieldIndex fieldIdx = kDataFieldStartingIdx;
 
   ArrayRef cooSegsRef = cooSegs;
@@ -951,7 +950,7 @@ Level mlir::sparse_tensor::SparseTensorEncodingAttr::getAoSCOOStart() const {
   return getLvlRank();
 }
 
-SmallVector<SparseTensorEncodingAttr::COOSegment>
+SmallVector<COOSegment>
 mlir::sparse_tensor::SparseTensorEncodingAttr::getCOOSegments() const {
   SmallVector<COOSegment> ret;
   if (getLvlRank() <= 1)