[mlir][sparse] introduce operations to query sparse tensor slice offset/strides at the given dimenion

[Peiming]

rebase.

[Peiming]

[mlir][sparse] extends storage specifier operations for slices.

[Peiming]

revert.

[aartbik]

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

301

can you add a bit more in the description other than an example
Having a simple illustrative example in terms of a tensor value helps too

[aartbik]

will these be implemented both along lib path as well as direct IR path?

[wrengr]

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

298

Do you really mean "dimension" here? or is this actually a "level"?

313

ditto

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

769

use getRankedTensorType

776

ditto

[Peiming]

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

298

I did mean dimension. Because slice should be applied on the user-faced shape (if you have high order mapping, the slice is not applied on storage level).

[wrengr]

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

298

Cool cool. I just wanted to make sure we got the right thing, so we don't need to change it later :)

[Peiming]

address comments.

[aartbik]

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

304

are still a work in progress

331

same

[Peiming]

address comments.

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

 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Traits.
 //===----------------------------------------------------------------------===//
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
+def IsSparseTensorSlicePred
+  : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self) && "
+          "  ::mlir::sparse_tensor::getSparseTensorEncoding($_self).isSlice()">;
+
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,
 // `RankedTensorOf`, `AnyRankedTensor`.
 
 class SparseTensorOf<list<Type> allowedTypes>
   : TensorOf<allowedTypes, [IsSparseTensorPred], "sparse tensor">;
 
+class SparseTensorSliceOf<list<Type> allowedTypes>
+  : TensorOf<allowedTypes, [IsSparseTensorSlicePred], "sparse tensor slice">;
+
 def AnySparseTensor : SparseTensorOf<[AnyType]>;
+def AnySparseTensorSlice : SparseTensorSliceOf<[AnyType]>;
 
 class RankedSparseTensorOf<list<Type> allowedTypes>
   : RankedTensorOf<allowedTypes, [IsSparseTensorPred], "ranked sparse tensor">;
 
 def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Sorting Algorithm Attribute.

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

let description = [{
     ```mlir
     %1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
   let hasVerifier = 1;
 }
 
+def SparseTensor_ToSliceOffsetOp : SparseTensor_Op<"slice.offset", [Pure]>,
+    Arguments<(ins AnySparseTensorSlice:$slice, IndexAttr:$dim)>,
+    Results<(outs Index:$offset)> {
+  let summary = "Extracts the offset of the sparse tensor slice at the given dimension";
+  let description = [{
+    Extracts the offset of the sparse tensor slice at the given dimension.
+
+    Currently, sparse tensor slices are still a work in progress, and only
+    works when runtime library is disabled (i.e., running sparse compiler
+    with `enable-runtime-library=false`).
+
+    Example:
+
+    ```mlir
+    %0 = tensor.extract_slice %s[%v1, %v2][64, 64][1, 1] : tensor<128x128xf64, #DCSR>
+                                                        to tensor<64x64xf64, #Slice>
+
+    %1 = sparse_tensor.slice.offset %0 at 0 : tensor<64x64xf64, #Slice>
+    %2 = sparse_tensor.slice.offset %0 at 1 : tensor<64x64xf64, #Slice>
+    // %1 = %v1
+    // %2 = %v2
+    ```
+  }];
+  let assemblyFormat = "$slice `at` $dim attr-dict `:` type($slice)";
+  let hasVerifier = 1;
+}
+
+def SparseTensor_ToSliceStrideOp : SparseTensor_Op<"slice.stride", [Pure]>,
+    Arguments<(ins AnySparseTensorSlice:$slice, IndexAttr:$dim)>,
+    Results<(outs Index:$stride)> {
+  let summary = "Extracts the stride of the sparse tensor slice at the given dimension";
+  let description = [{
+    Extracts the stride of the sparse tensor slice at the given dimension.
+
+    Currently, sparse tensor slices are still a work in progress, and only
+    works when runtime library is disabled (i.e., running sparse compiler
+    with `enable-runtime-library=false`).
+
+    Example:
+
+    ```mlir
+    %0 = tensor.extract_slice %s[%v1, %v2][64, 64][%s1, %s2] : tensor<128x128xf64, #DCSR>
+                                                            to tensor<64x64xf64, #Slice>
+
+    %1 = sparse_tensor.slice.stride %0 at 0 : tensor<64x64xf64, #Slice>
+    %2 = sparse_tensor.slice.stride %0 at 1 : tensor<64x64xf64, #Slice>
+    // %1 = %s1
+    // %2 = %s2
+
+    ```
+  }];
+  let assemblyFormat = "$slice `at` $dim attr-dict `:` type($slice)";
+  let hasVerifier = 1;
+}
+
 def SparseTensor_StorageSpecifierInitOp : SparseTensor_Op<"storage_specifier.init", [Pure]>,
     Results<(outs SparseTensorStorageSpecifier:$result)> {
   let summary = "";
   let description = [{
     Returns an initial storage specifier value. A storage specifier value holds
     the sizes for tensor dimensions, pointer arrays, index arrays, and the value array.
 
     Example:

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

 LogicalResult ToValuesOp::verify() {
   auto ttp = getRankedTensorType(getTensor());
   auto mtp = getMemRefType(getResult());
   if (ttp.getElementType() != mtp.getElementType())
     return emitError("unexpected mismatch in element types");
   return success();
 }
 
+LogicalResult ToSliceOffsetOp::verify() {
+  auto rank = getRankedTensorType(getSlice()).getRank();
+  if (rank <= getDim().getSExtValue() || getDim().getSExtValue() < 0)
+    return emitError("requested dimension out of bound");
+  return success();
+}
+
+LogicalResult ToSliceStrideOp::verify() {
+  auto rank = getRankedTensorType(getSlice()).getRank();
+  if (rank <= getDim().getSExtValue() || getDim().getSExtValue() < 0)
+    return emitError("requested dimension out of bound");
+  return success();
+}
+
 LogicalResult GetStorageSpecifierOp::verify() {
   RETURN_FAILURE_IF_FAILED(verifySparsifierGetterSetter(
       getSpecifierKind(), getDim(), getSpecifier(), getOperation()))
   // Checks the result type
   if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=
       getResult().getType()) {
     return emitError(
         "type mismatch between requested specifier field and result value");

mlir/test/Dialect/SparseTensor/invalid.mlir

 func.func @mismatch_values_types(%arg0: tensor<?xf64, #SparseVector>) -> memref<?xf32> {
   // expected-error@+1 {{unexpected mismatch in element types}}
   %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf32>
   return %0 : memref<?xf32>
 }
 
 // -----
 
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+func.func @sparse_slice_offset(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  // expected-error@+1 {{requested dimension out of bound}}
+  %0 = sparse_tensor.slice.offset %arg0 at 2 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+func.func @sparse_slice_stride(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  // expected-error@+1 {{requested dimension out of bound}}
+  %0 = sparse_tensor.slice.stride %arg0 at 2 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{redundant level argument for querying value memory size}}
   %0 = sparse_tensor.storage_specifier.get %arg0 val_mem_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: return %[[T]] : memref<?xf64>
 func.func @sparse_values(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // -----
 
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func @sparse_slice_offset(
+//  CHECK-SAME: %[[A:.*]]: tensor<2x8xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.slice.offset %[[A]] at 1 : tensor<2x8xf64, #{{.*}}>
+//       CHECK: return %[[T]] : index
+func.func @sparse_slice_offset(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  %0 = sparse_tensor.slice.offset %arg0 at 1 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
+#CSR_SLICE = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+// CHECK-LABEL: func @sparse_slice_stride(
+//  CHECK-SAME: %[[A:.*]]: tensor<2x8xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.slice.stride %[[A]] at 1 : tensor<2x8xf64, #{{.*}}>
+//       CHECK: return %[[T]] : index
+func.func @sparse_slice_stride(%arg0: tensor<2x8xf64, #CSR_SLICE>) -> index {
+  %0 = sparse_tensor.slice.stride %arg0 at 1 : tensor<2x8xf64, #CSR_SLICE>
+  return %0 : index
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_metadata_init(
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#SparseVector> {
   %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>