faa00c1313
A previous revision implemented expand/collapse reshaping between dense and sparse tensors for sparse2dense and dense2sparse since those could use the "cheap" view reshape on the already materialized dense tensor (at either the input or output side), and do some reshuffling from or to sparse. The dense2dense case, as always, is handled with a "cheap" view change. This revision implements the sparse2sparse cases. Lacking any "view" support on sparse tensors this operation necessarily has to perform data reshuffling on both ends. Tracker for improving this: https://github.com/llvm/llvm-project/issues/56477 Reviewed By: bixia Differential Revision: https://reviews.llvm.org/D129416
100 lines
5.3 KiB
MLIR
Executable File
100 lines
5.3 KiB
MLIR
Executable File
// RUN: mlir-opt %s -sparsification | FileCheck %s
|
|
|
|
#SparseVector = #sparse_tensor.encoding<{
|
|
dimLevelType = ["compressed"]
|
|
}>
|
|
|
|
#SparseMatrix = #sparse_tensor.encoding<{
|
|
dimLevelType = ["compressed", "compressed"]
|
|
}>
|
|
|
|
// CHECK-LABEL: func.func @expand_dense(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<12xf64>) -> tensor<3x4xf64> {
|
|
// CHECK: %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
|
|
// CHECK: return %[[E]] : tensor<3x4xf64>
|
|
// CHECK: }
|
|
func.func @expand_dense(%arg0: tensor<12xf64>) -> tensor<3x4xf64> {
|
|
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64>
|
|
return %0 : tensor<3x4xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @expand_from_sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<3x4xf64> {
|
|
// CHECK: %[[C:.*]] = sparse_tensor.convert %[[A]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<12xf64>
|
|
// CHECK: %[[E:.*]] = tensor.expand_shape %[[C]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
|
|
// CHECK: return %[[E]] : tensor<3x4xf64>
|
|
// CHECK: }
|
|
func.func @expand_from_sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> {
|
|
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64>
|
|
return %0 : tensor<3x4xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @expand_to_sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<12xf64>) -> tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
|
|
// CHECK: %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
|
|
// CHECK: %[[C:.*]] = sparse_tensor.convert %[[E]] : tensor<3x4xf64> to tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: return %[[C]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: }
|
|
func.func @expand_to_sparse(%arg0: tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> {
|
|
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix>
|
|
return %0 : tensor<3x4xf64, #SparseMatrix>
|
|
}
|
|
|
|
//
|
|
// Not rewritten, needs conversion.
|
|
//
|
|
// CHECK-LABEL: func.func @expand_sparse2sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
|
|
// CHECK: %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: return %[[E]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: }
|
|
func.func @expand_sparse2sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> {
|
|
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix>
|
|
return %0 : tensor<3x4xf64, #SparseMatrix>
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @collapse_dense(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<3x4xf64>) -> tensor<12xf64> {
|
|
// CHECK: %[[R:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
|
|
// CHECK: return %[[R]] : tensor<12xf64>
|
|
// CHECK: }
|
|
func.func @collapse_dense(%arg0: tensor<3x4xf64>) -> tensor<12xf64> {
|
|
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64>
|
|
return %0 : tensor<12xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @collapse_from_sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<12xf64> {
|
|
// CHECK: %[[C:.*]] = sparse_tensor.convert %[[A]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<3x4xf64>
|
|
// CHECK: %[[R:.*]] = tensor.collapse_shape %[[C]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
|
|
// CHECK: return %[[R]] : tensor<12xf64>
|
|
// CHECK: }
|
|
func.func @collapse_from_sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> {
|
|
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64>
|
|
return %0 : tensor<12xf64>
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @collapse_to_sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<3x4xf64>) -> tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> {
|
|
// CHECK: %[[R:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
|
|
// CHECK: %[[C:.*]] = sparse_tensor.convert %[[R]] : tensor<12xf64> to tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: return %[[C]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: }
|
|
func.func @collapse_to_sparse(%arg0: tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> {
|
|
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64, #SparseVector>
|
|
return %0 : tensor<12xf64, #SparseVector>
|
|
}
|
|
|
|
//
|
|
// Not rewritten, needs conversion.
|
|
//
|
|
// CHECK-LABEL: func.func @collapse_sparse2sparse(
|
|
// CHECK-SAME: %[[A:.*]]: tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> {
|
|
// CHECK: %[[C:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: return %[[C]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
|
|
// CHECK: }
|
|
func.func @collapse_sparse2sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> {
|
|
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64, #SparseVector>
|
|
return %0 : tensor<12xf64, #SparseVector>
|
|
}
|