[mlir][sparse] Add lowering rules for sparse_tensor.storage Op

[aartbik]

mlir/test/Dialect/SparseTensor/codegen.mlir

293

minor nitpick, these blocks of code are going to be almost literal repeats, except for the tuple part.
perhaps we should just check for the tuple parts in the STORAGE part

at the very least, let's use the same names for all variables that are shared?

[Peiming]

address comments from Aart

[Peiming]

mlir/test/Dialect/SparseTensor/codegen.mlir

293

How about the current way? I found that you can define multiple prefixes so that the common part can be checked by shared "CHECK" prefix, while the different parts can be checked by their own prefix.

[aartbik]

This is beautiful! Thanks.

[Peiming]

rebase

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

for (auto operand : operands) {
     else
       flattened.push_back(operand);
   }
 }
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
+/// Sparse tensor storage conversion rule for sparse_tensor::storage.
+class SparseStorageConversion : public OpConversionPattern<StorageOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(StorageOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Simply convert it to a unrealize_conversion_cast.
+    // We should guarantee that all uses of sparse_tensor.storage op will
+    // be eventually eliminated by accessing the flattened SSA values directly.
+    rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
+        op, TypeRange{op.getType()}, adaptor.getInputs());
+    return success();
+  }
+};
+
 /// Sparse tensor storage conversion rule for sparse_tensor::storage_get.
 class SparseStorageGetConverter : public OpConversionPattern<StorageGetOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(StorageGetOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto castOp =

 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required
 /// to expand compounded sparse tensor tuples.
 void mlir::populateSparseTensorStorageExpansionPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns) {
-  patterns.add<SparseStorageGetConverter, SparseStorageSetConverter,
-               SparseStorageReturnConverter, SparseStorageCallConverter>(
-      typeConverter, patterns.getContext());
+  patterns.add<SparseStorageConversion, SparseStorageGetConverter,
+               SparseStorageSetConverter, SparseStorageReturnConverter,
+               SparseStorageCallConverter>(typeConverter,
+                                           patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

-// RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize --cse | FileCheck %s --check-prefix=CHECK-CODEGEN
-// FIXME:
-// R_U_N: mlir-opt %s --sparse-tensor-codegen --sparse-tensor-storage-expansion --canonicalize --cse | FileCheck %s --check-prefix=CHECK-STORAGE
+// RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize --cse | FileCheck %s --check-prefixes=CHECK,CHECK-CODEGEN
+// RUN: mlir-opt %s --sparse-tensor-codegen --sparse-tensor-storage-expansion --canonicalize --cse | FileCheck %s --check-prefixes=CHECK,CHECK-STORAGE
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed" ],
   indexBitWidth = 64,
   pointerBitWidth = 32
 }>
 
 #Dense2D = #sparse_tensor.encoding<{

 //       CHECK-STORAGE: memref.dealloc %[[A2]] : memref<?xi64>
 //       CHECK-STORAGE: memref.dealloc %[[A3]] : memref<?xf64>
 //       CHECK-STORAGE: return
 func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {
   bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>
   return
 }
 
-// CHECK-CODEGEN-LABEL: func @sparse_alloc_csc(
-//  CHECK-CODEGEN-SAME: %[[A:.*]]: index)
-//   CHECK-CODEGEN-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-CODEGEN-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-CODEGEN-DAG: %[[C10:.*]] = arith.constant 10 : index
-//      CHECK-CODEGEN:  %[[T0:.*]] = memref.alloc() : memref<2xindex>
-//      CHECK-CODEGEN:  memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
-//      CHECK-CODEGEN:  memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
-//      CHECK-CODEGEN:  %[[T1:.*]] = memref.alloc() : memref<1xindex>
-//      CHECK-CODEGEN:  %[[T2:.*]] = memref.cast %[[T1]] : memref<1xindex> to memref<?xindex>
-//      CHECK-CODEGEN:  %[[T3:.*]] = memref.alloc() : memref<1xindex>
-//      CHECK-CODEGEN:  %[[T4:.*]] = memref.cast %[[T3]] : memref<1xindex> to memref<?xindex>
-//      CHECK-CODEGEN:  %[[T5:.*]] = memref.alloc() : memref<1xf64>
-//      CHECK-CODEGEN:  %[[T6:.*]] = memref.cast %[[T5]] : memref<1xf64> to memref<?xf64>
-//      CHECK-CODEGEN:  %[[T:.*]] = sparse_tensor.storage(%[[T0]], %[[T2]], %[[T4]], %[[T6]])
-//      CHECK-CODEGEN:  return %[[T]] : tuple<memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>>
+//        CHECK-LABEL: func @sparse_alloc_csc(
+//         CHECK-SAME: %[[A:.*]]: index) ->
+// CHECK-CODEGEN-SAME: tuple<memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>>
+// CHECK-STORAGE-SAME: memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//          CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//          CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//          CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//              CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>
+//              CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
+//              CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
+//              CHECK: %[[T1:.*]] = memref.alloc() : memref<1xindex>
+//              CHECK: %[[T2:.*]] = memref.cast %[[T1]] : memref<1xindex> to memref<?xindex>
+//              CHECK: %[[T3:.*]] = memref.alloc() : memref<1xindex>
+//              CHECK: %[[T4:.*]] = memref.cast %[[T3]] : memref<1xindex> to memref<?xindex>
+//              CHECK: %[[T5:.*]] = memref.alloc() : memref<1xf64>
+//              CHECK: %[[T6:.*]] = memref.cast %[[T5]] : memref<1xf64> to memref<?xf64>
+//      CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[T0]], %[[T2]], %[[T4]], %[[T6]])
+//      CHECK-CODEGEN: return %[[T]]
+//      CHECK-STORAGE: return %[[T0]], %[[T2]], %[[T4]], %[[T6]] 
 func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {
   %0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>
   %1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>
   return %1 : tensor<10x?xf64, #CSC>
 }
 
-// CHECK-CODEGEN-LABEL: func @sparse_alloc_3d() -> tuple<memref<3xindex>, memref<?xf64>>
-//   CHECK-CODEGEN-DAG: %[[C0:.*]] = arith.constant 0 : index
-//   CHECK-CODEGEN-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-CODEGEN-DAG: %[[C2:.*]] = arith.constant 2 : index
-//   CHECK-CODEGEN-DAG: %[[C10:.*]] = arith.constant 10 : index
-//   CHECK-CODEGEN-DAG: %[[C20:.*]] = arith.constant 20 : index
-//   CHECK-CODEGEN-DAG: %[[C30:.*]] = arith.constant 30 : index
-//       CHECK-CODEGEN: %[[A0:.*]] = memref.alloc() : memref<3xindex>
-//       CHECK-CODEGEN: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
-//       CHECK-CODEGEN: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
-//       CHECK-CODEGEN: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
-//       CHECK-CODEGEN: %[[A:.*]] = memref.alloc() : memref<6000xf64>
-//       CHECK-CODEGEN: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
-//       CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]])
-//       CHECK-CODEGEN: return %[[T]] : tuple<memref<3xindex>, memref<?xf64>>
+//        CHECK-LABEL: func @sparse_alloc_3d() ->
+// CHECK-CODEGEN-SAME: tuple<memref<3xindex>, memref<?xf64>>
+// CHECK-STORAGE-SAME: memref<3xindex>, memref<?xf64>
+//          CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//          CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//          CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//          CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//          CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
+//          CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index
+//              CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>
+//              CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
+//              CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
+//              CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
+//              CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>
+//              CHECK: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
+//      CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]])
+//      CHECK-CODEGEN: return %[[T]] : tuple<memref<3xindex>, memref<?xf64>>
+//      CHECK-STORAGE: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf64>
 func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
   %0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>
   %1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>
   return %1 : tensor<10x20x30xf64, #Dense3D>
 }

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

 // RUN: mlir-opt %s -sparse-tensor-storage-expansion -cse | FileCheck %s
 
 // CHECK-LABEL:  func @sparse_storage_expand(
 // CHECK-SAME:     %[[TMP_arg0:.*0]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg1:.*1]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg2:.*]]: f64
 // CHECK           return %[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]
 func.func @sparse_storage_expand(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>)
                                      -> tuple<memref<?xf64>, memref<?xf64>, f64> {
   return %arg0 : tuple<memref<?xf64>, memref<?xf64>, f64>
 }
 
 // CHECK-LABEL:  func @call_sparse_storage_expand(
 // CHECK-SAME:     %[[TMP_arg0:.*0]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg1:.*1]]: memref<?xf64>,
-// CHECK-SAME:     %[[TMP_arg2:.*]]: f64) 
-// CHECK:          %[[TMP_0:.*]]:3 = call @sparse_storage_expand(%[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]) 
+// CHECK-SAME:     %[[TMP_arg2:.*]]: f64)
+// CHECK:          %[[TMP_0:.*]]:3 = call @sparse_storage_expand(%[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]])
 // CHECK:          return %[[TMP_0]]#0, %[[TMP_0]]#1, %[[TMP_0]]#2 : memref<?xf64>, memref<?xf64>, f64
 func.func @call_sparse_storage_expand(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>)
                                           -> tuple<memref<?xf64>, memref<?xf64>, f64> {
   %1 = call @sparse_storage_expand(%arg0) : (tuple<memref<?xf64>, memref<?xf64>, f64>) ->
                                              tuple<memref<?xf64>, memref<?xf64>, f64>
   return %1 : tuple<memref<?xf64>, memref<?xf64>, f64>
 }
 
+// CHECK-LABEL: func @sparse_storage(
+// CHECK-SAME:    %[[TMP_arg0:.*0]]: memref<?xf64>,
+// CHECK-SAME:    %[[TMP_arg1:.*1]]: memref<?xf64>,
+// CHECK-SAME:    %[[TMP_arg2:.*2]]: memref<?xf64>)
+// CHECK:         return %[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]
+func.func @sparse_storage(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: memref<?xf64>)
+                        -> tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>> {
+  %1 = sparse_tensor.storage(%arg0, %arg1, %arg2) : memref<?xf64>, memref<?xf64>, memref<?xf64> to tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>>
+  return %1 : tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>>
+}
+
 // CHECK-LABEL:  func @sparse_storage_get(
 // CHECK-SAME:     %[[TMP_arg0:.*0]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg1:.*1]]: memref<?xf64>,
-// CHECK-SAME:     %[[TMP_arg2:.*]]: f64) 
+// CHECK-SAME:     %[[TMP_arg2:.*]]: f64)
 // CHECK:          return %[[TMP_arg0]] : memref<?xf64>
 func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
   %0 = sparse_tensor.storage_get %arg0[0]
        : tuple<memref<?xf64>, memref<?xf64>, f64> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // CHECK-LABEL:  func @sparse_storage_set(
 // CHECK-SAME:     %[[TMP_arg0:.*0]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg1:.*1]]: memref<?xf64>,
 // CHECK-SAME:     %[[TMP_arg2:.*]]: f64,
-// CHECK-SAME:     %[[TMP_arg3:.*]]: memref<?xf64>) 
+// CHECK-SAME:     %[[TMP_arg3:.*]]: memref<?xf64>)
 // CHECK:          return %[[TMP_arg3]], %[[TMP_arg1]], %[[TMP_arg2]] : memref<?xf64>, memref<?xf64>, f64
 func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>,
                               %arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
   %0 = sparse_tensor.storage_set %arg0[0], %arg1
        : tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
          tuple<memref<?xf64>, memref<?xf64>, f64>
   return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
 }