From 662ee93cbbaad27e8dfecfdd4731dc4b168658b4 Mon Sep 17 00:00:00 2001 From: Peiming Liu Date: Fri, 2 Sep 2022 17:31:54 +0000 Subject: [PATCH] [mlir][sparse] extend codegen test cases with an additional step after storage expansion Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D133219 --- mlir/test/Dialect/SparseTensor/codegen.mlir | 202 +++++++++++++++++++++------- 1 file changed, 151 insertions(+), 51 deletions(-) diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir index 6d95be8..4da1d0b 100644 --- a/mlir/test/Dialect/SparseTensor/codegen.mlir +++ b/mlir/test/Dialect/SparseTensor/codegen.mlir @@ -1,4 +1,6 @@ -// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse | FileCheck %s +// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse | FileCheck %s --check-prefix=CHECK-CODEGEN +// RUN: mlir-opt %s --sparse-tensor-codegen --sparse-tensor-storage-expansion --canonicalize --cse | FileCheck %s --check-prefix=CHECK-STORAGE + #SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], @@ -35,49 +37,96 @@ dimOrdering = affine_map<(i, j, k) -> (k, i, j)> }> -// CHECK-LABEL: func @sparse_nop( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref>) -// CHECK: return %[[A]] : tuple, memref, memref, memref> +// CHECK-CODEGEN-LABEL: func @sparse_nop( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// CHECK-CODEGEN: return %[[A]] : tuple, memref, memref, memref> +// +// CHECK-STORAGE-LABEL: func @sparse_nop( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<1xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref) +// CHECK-STORAGE: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref, memref, memref func.func @sparse_nop(%arg0: tensor) -> tensor { return %arg0 : tensor } -// CHECK-LABEL: func @sparse_nop_cast( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref>) -// CHECK: return %[[A]] : tuple, memref, memref, memref> +// CHECK-CODEGEN-LABEL: func @sparse_nop_cast( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// CHECK-CODEGEN: return %[[A]] : tuple, memref, memref, memref> +// +// CHECK-STORAGE-LABEL: func @sparse_nop_cast( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<1xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref) +// CHECK-STORAGE: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref, memref, memref func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor { %0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor return %0 : tensor } -// CHECK-LABEL: func @sparse_nop_cast_3d( -// CHECK-SAME: %[[A:.*]]: tuple, memref>) -// CHECK: return %[[A]] : tuple, memref> +// CHECK-CODEGEN-LABEL: func @sparse_nop_cast_3d( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref>) +// CHECK-CODEGEN: return %[[A]] : tuple, memref> +// +// CHECK-STORAGE-LABEL: func @sparse_nop_cast_3d( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref) +// CHECK-STORAGE: return %[[A0]], %[[A1]] : memref<3xindex>, memref func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor { %0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor return %0 : tensor } -// CHECK-LABEL: func @sparse_dense_2d( -// CHECK-SAME: %[[A:.*]]: tuple, memref>) +// CHECK-CODEGEN-LABEL: func @sparse_dense_2d( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref>) +// +// CHECK-STORAGE-LABEL: func @sparse_dense_2d( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref) { +// CHECK-STORAGE: return func.func @sparse_dense_2d(%arg0: tensor) { return } -// CHECK-LABEL: func @sparse_row( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// CHECK-CODEGEN-LABEL: func @sparse_row( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// +// CHECK-STORAGE-LABEL: func @sparse_row( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref) { +// CHECK-STORAGE: return func.func @sparse_row(%arg0: tensor) { return } -// CHECK-LABEL: func @sparse_csr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// CHECK-CODEGEN-LABEL: func @sparse_csr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// +// CHECK-STORAGE-LABEL: func @sparse_csr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref) { +// CHECK-STORAGE: return func.func @sparse_csr(%arg0: tensor) { return } -// CHECK-LABEL: func @sparse_dcsr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) +// CHECK-CODEGEN-LABEL: func @sparse_dcsr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) +// +// CHECK-STORAGE-LABEL: func @sparse_dcsr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref, +// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref, +// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref) { +// CHECK-STORAGE: return func.func @sparse_dcsr(%arg0: tensor) { return } @@ -86,10 +135,16 @@ func.func @sparse_dcsr(%arg0: tensor) { // Querying for dimension 1 in the tensor type can immediately // fold using the original static dimension sizes. // -// CHECK-LABEL: func @sparse_dense_3d( -// CHECK-SAME: %[[A:.*]]: tuple, memref>) -// CHECK: %[[C:.*]] = arith.constant 20 : index -// CHECK: return %[[C]] : index +// CHECK-CODEGEN-LABEL: func @sparse_dense_3d( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref>) +// CHECK-CODEGEN: %[[C:.*]] = arith.constant 20 : index +// CHECK-CODEGEN: return %[[C]] : index +// +// CHECK-STORAGE-LABEL: func @sparse_dense_3d( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref) +// CHECK-STORAGE: %[[C:.*]] = arith.constant 20 : index +// CHECK-STORAGE: return %[[C]] : index func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index { %c = arith.constant 1 : index %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D> @@ -101,58 +156,103 @@ func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index { // into querying for dimension 2 in the stored sparse tensor scheme, // since the latter honors the dimOrdering. // -// CHECK-LABEL: func @sparse_dense_3d_dyn( -// CHECK-SAME: %[[A:.*]]: tuple, memref>) -// CHECK: %[[C:.*]] = arith.constant 2 : index -// CHECK: %[[F:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple, memref> to memref<3xindex> -// CHECK: %[[L:.*]] = memref.load %[[F]][%[[C]]] : memref<3xindex> -// CHECK: return %[[L]] : index +// CHECK-CODEGEN-LABEL: func @sparse_dense_3d_dyn( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref>) +// CHECK-CODEGEN: %[[C:.*]] = arith.constant 2 : index +// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple, memref> to memref<3xindex> +// CHECK-CODEGEN: %[[L:.*]] = memref.load %[[F]][%[[C]]] : memref<3xindex> +// CHECK-CODEGEN: return %[[L]] : index +// +// CHECK-STORAGE-LABEL: func @sparse_dense_3d_dyn( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref) +// CHECK-STORAGE: %[[C:.*]] = arith.constant 2 : index +// CHECK-STORAGE: %[[L:.*]] = memref.load %[[A0]][%[[C]]] : memref<3xindex> +// CHECK-STORAGE: return %[[L]] : index func.func @sparse_dense_3d_dyn(%arg0: tensor) -> index { %c = arith.constant 1 : index %0 = tensor.dim %arg0, %c : tensor return %0 : index } -// CHECK-LABEL: func @sparse_pointers_dcsr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) -// CHECK: %[[F:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple, memref, memref, memref, memref, memref> to memref -// CHECK: return %[[F]] : memref +// CHECK-CODEGEN-LABEL: func @sparse_pointers_dcsr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) +// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple, memref, memref, memref, memref, memref> to memref +// CHECK-CODEGEN: return %[[F]] : memref +// +// CHECK-STORAGE-LABEL: func @sparse_pointers_dcsr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref, +// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref, +// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref) +// CHECK-STORAGE: return %[[A3]] : memref func.func @sparse_pointers_dcsr(%arg0: tensor) -> memref { %c = arith.constant 1 : index %0 = sparse_tensor.pointers %arg0, %c : tensor to memref return %0 : memref } -// CHECK-LABEL: func @sparse_indices_dcsr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) -// CHECK: %[[F:.*]] = sparse_tensor.storage_get %[[A]][4] : tuple, memref, memref, memref, memref, memref> to memref -// CHECK: return %[[F]] : memref +// CHECK-CODEGEN-LABEL: func @sparse_indices_dcsr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) +// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][4] : tuple, memref, memref, memref, memref, memref> to memref +// CHECK-CODEGEN: return %[[F]] : memref +// +// CHECK-STORAGE-LABEL: func @sparse_indices_dcsr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref, +// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref, +// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref) +// CHECK-STORAGE: return %[[A4]] : memref func.func @sparse_indices_dcsr(%arg0: tensor) -> memref { %c = arith.constant 1 : index %0 = sparse_tensor.indices %arg0, %c : tensor to memref return %0 : memref } -// CHECK-LABEL: func @sparse_values_dcsr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) -// CHECK: %[[F:.*]] = sparse_tensor.storage_get %[[A]][5] : tuple, memref, memref, memref, memref, memref> to memref -// CHECK: return %[[F]] : memref +// CHECK-CODEGEN-LABEL: func @sparse_values_dcsr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref, memref, memref>) +// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][5] : tuple, memref, memref, memref, memref, memref> to memref +// CHECK-CODEGEN: return %[[F]] : memref +// +// CHECK-STORAGE-LABEL: func @sparse_values_dcsr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref, +// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref, +// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref) +// CHECK-STORAGE: return %[[A5]] : memref func.func @sparse_values_dcsr(%arg0: tensor) -> memref { %0 = sparse_tensor.values %arg0 : tensor to memref return %0 : memref } -// CHECK-LABEL: func @sparse_dealloc_csr( -// CHECK-SAME: %[[A:.*]]: tuple, memref, memref, memref>) -// CHECK: %[[F0:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple, memref, memref, memref> to memref<2xindex> -// CHECK: memref.dealloc %[[F0]] : memref<2xindex> -// CHECK: %[[F1:.*]] = sparse_tensor.storage_get %[[A]][1] : tuple, memref, memref, memref> to memref -// CHECK: memref.dealloc %[[F1]] : memref -// CHECK: %[[F2:.*]] = sparse_tensor.storage_get %[[A]][2] : tuple, memref, memref, memref> to memref -// CHECK: memref.dealloc %[[F2]] : memref -// CHECK: %[[F3:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple, memref, memref, memref> to memref -// CHECK: memref.dealloc %[[F3]] : memref -// CHECK: return +// CHECK-CODEGEN-LABEL: func @sparse_dealloc_csr( +// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple, memref, memref, memref>) +// CHECK-CODEGEN: %[[F0:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple, memref, memref, memref> to memref<2xindex> +// CHECK-CODEGEN: memref.dealloc %[[F0]] : memref<2xindex> +// CHECK-CODEGEN: %[[F1:.*]] = sparse_tensor.storage_get %[[A]][1] : tuple, memref, memref, memref> to memref +// CHECK-CODEGEN: memref.dealloc %[[F1]] : memref +// CHECK-CODEGEN: %[[F2:.*]] = sparse_tensor.storage_get %[[A]][2] : tuple, memref, memref, memref> to memref +// CHECK-CODEGEN: memref.dealloc %[[F2]] : memref +// CHECK-CODEGEN: %[[F3:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple, memref, memref, memref> to memref +// CHECK-CODEGEN: memref.dealloc %[[F3]] : memref +// CHECK-CODEGEN: return +// +// CHECK-STORAGE-LABEL: func @sparse_dealloc_csr( +// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>, +// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref, +// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref, +// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref) { +// CHECK-STORAGE: memref.dealloc %[[A0]] : memref<2xindex> +// CHECK-STORAGE: memref.dealloc %[[A1]] : memref +// CHECK-STORAGE: memref.dealloc %[[A2]] : memref +// CHECK-STORAGE: memref.dealloc %[[A3]] : memref +// CHECK-STORAGE: return func.func @sparse_dealloc_csr(%arg0: tensor) { bufferization.dealloc_tensor %arg0 : tensor return -- 2.7.4