[mlir][sparse] add folder to sparse_tensor.storage.get operation.

[aartbik]

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

498

can you make this a for loop with the init, test and update all in one line?

[Peiming]

use for loop

[Peiming]

code simplication.

[Peiming]

more consices

[aartbik]

Beautiful, ship it!

[wrengr]

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

497

Why not use x.getSpecifier().getDefiningOp... since that's how it's called in both of the use sites?

[Peiming]

make a template

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

let description = [{
     %0 = sparse_tensor.storage_specifier.get %arg0 idx_mem_sz at 0
          : !sparse_tensor.storage_specifier<#COO> to i64
     ```
   }];
 
   let assemblyFormat = "$specifier $specifierKind (`at` $dim^)? attr-dict `:` "
                        "qualified(type($specifier)) `to` type($result)";
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def SparseTensor_SetStorageSpecifierOp : SparseTensor_Op<"storage_specifier.set",
     [Pure, AllTypesMatch<["result", "specifier"]>]>,
     Arguments<(ins SparseTensorStorageSpecifier:$specifier,
                    SparseTensorStorageSpecifierKindAttr:$specifierKind,
                    OptionalAttr<IndexAttr>:$dim,
                    AnyType:$value)>,

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

LogicalResult GetStorageSpecifierOp::verify() {
   if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=
       getResult().getType()) {
     return emitError(
         "type mismatch between requested specifier field and result value");
   }
   return success();
 }
 
+template <typename SpecifierOp>
+static SetStorageSpecifierOp getSpecifierSetDef(SpecifierOp op) {
+  return op.getSpecifier().template getDefiningOp<SetStorageSpecifierOp>();
+}
+
+OpFoldResult GetStorageSpecifierOp::fold(ArrayRef<Attribute> operands) {
+  StorageSpecifierKind kind = getSpecifierKind();
+  Optional<APInt> dim = getDim();
+  for (auto op = getSpecifierSetDef(*this); op; op = getSpecifierSetDef(op))
+    if (kind == op.getSpecifierKind() && dim == op.getDim())
+      return op.getValue();
+  return {};
+}
+
 LogicalResult SetStorageSpecifierOp::verify() {
   if (failed(verifySparsifierGetterSetter(getSpecifierKind(), getDim(),
                                           getSpecifier(), getOperation()))) {
     return failure();
   }
 
   // Checks the input type
   if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=

mlir/test/Dialect/SparseTensor/fold.mlir

 //   CHECK-NOT: sparse_tensor.values
 //       CHECK: return
 func.func @sparse_dce_getters(%arg0: tensor<64xf32, #SparseVector>) {
   %0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<64xf32, #SparseVector> to memref<?xindex>
   %1 = sparse_tensor.indices %arg0 { dimension = 0 : index } : tensor<64xf32, #SparseVector> to memref<?xindex>
   %2 = sparse_tensor.values %arg0 : tensor<64xf32, #SparseVector> to memref<?xf32>
   return
 }
+
 // CHECK-LABEL: func @sparse_concat_dce(
 //   CHECK-NOT: sparse_tensor.concatenate
 //       CHECK: return
 func.func @sparse_concat_dce(%arg0: tensor<2xf64, #SparseVector>,
                              %arg1: tensor<3xf64, #SparseVector>,
                              %arg2: tensor<4xf64, #SparseVector>) {
   %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2xf64, #SparseVector>,
          tensor<3xf64, #SparseVector>,
          tensor<4xf64, #SparseVector> to tensor<9xf64, #SparseVector>
   return
 }
 
+// CHECK-LABEL: func @sparse_get_specifier_dce_fold(
+//  CHECK-SAME:  %[[A0:.*]]: !sparse_tensor.storage_specifier
+//  CHECK-SAME:  %[[A1:.*]]: i64,
+//  CHECK-SAME:  %[[A2:.*]]: i64)
+//   CHECK-NOT:  sparse_tensor.storage_specifier.set
+//   CHECK-NOT:  sparse_tensor.storage_specifier.get
+//       CHECK:  return %[[A1]]
+func.func @sparse_get_specifier_dce_fold(%arg0: !sparse_tensor.storage_specifier<#SparseVector>, %arg1: i64, %arg2: i64) -> i64 {
+  %0 = sparse_tensor.storage_specifier.set %arg0 dim_sz at 0 with %arg1
+       : i64, !sparse_tensor.storage_specifier<#SparseVector>
+  %1 = sparse_tensor.storage_specifier.set %0 ptr_mem_sz at 0 with %arg2
+       : i64, !sparse_tensor.storage_specifier<#SparseVector>
+  %2 = sparse_tensor.storage_specifier.get %1 dim_sz at 0
+       : !sparse_tensor.storage_specifier<#SparseVector> to i64
+  return %2 : i64
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   return %0 : i64
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_set_md(
-//  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>, 
+//  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>,
 //  CHECK-SAME: %[[I:.*]]: i64)
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.set %[[A]] dim_sz at 0 with %[[I]]
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>, %arg1: i64)
           -> !sparse_tensor.storage_specifier<#SparseVector> {
   %0 = sparse_tensor.storage_specifier.set %arg0 dim_sz at 0 with %arg1
        : i64, !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>

 }
 
 // -----
 
 func.func @sparse_sort_coo_stable(%arg0: index, %arg1: memref<?xi64>, %arg2: memref<?xf32>) -> (memref<?xi64>, memref<?xf32>) {
   sparse_tensor.sort_coo stable %arg0, %arg1 jointly %arg2 { nx=2 : index, ny=1 : index}: memref<?xi64> jointly memref<?xf32>
   return %arg1, %arg2 : memref<?xi64>, memref<?xf32>
 }
-