Skip to content

Commit

Permalink
[mlir][sparse] fix sparse tests that uses reshape operations. (#90637)
Browse files Browse the repository at this point in the history 
Due to generalization introduced in
#90040
  • Loading branch information
@PeimingLiu
PeimingLiu committed Apr 30, 2024
1 parent 41f9c78 commit 7cbaaed
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 16 deletions.
8 changes: 4 additions & 4 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/reshape_dot.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ module {
%cst = arith.constant 0.000000e+00 : f32
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
%2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32>, tensor<6x6xf32, #COO_2D>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] output_shape [5,2,3]: tensor<5x6xf32> into tensor<5x2x3xf32>
%ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>

// Note: tensor.collapse_shape is a metadata-only operation on dense tensors
Expand All @@ -60,7 +60,7 @@ module {
%cst = arith.constant 0.000000e+00 : f32
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
%2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32, #COO_2D>, tensor<6x6xf32, #COO_2D>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] output_shape [5,2,3]: tensor<5x6xf32> into tensor<5x2x3xf32>
%ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>

// Note: tensor.collapse_shape is a metadata-only operation on dense tensors
Expand All @@ -76,7 +76,7 @@ module {
%cst = arith.constant 0.000000e+00 : f32
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
%2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32>, tensor<6x6xf32>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] output_shape [5,2,3]: tensor<5x6xf32> into tensor<5x2x3xf32>
%ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>
return %ret1 : tensor<?x?x?xf32>
}
Expand All @@ -88,7 +88,7 @@ module {
%cst = arith.constant 0.000000e+00 : f32
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
%2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32, #COO_2D>, tensor<6x6xf32, #COO_2D>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
%expanded = tensor.expand_shape %2 [[0], [1, 2]] output_shape [5,2,3]: tensor<5x6xf32> into tensor<5x2x3xf32>
%ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>

// Note: tensor.collapse_shape is a metadata-only operation on dense tensors
Expand Down
48 changes: 36 additions & 12 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand_shape.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,62 +53,86 @@
module {

func.func @expand_dense(%arg0: tensor<12xf64>) -> tensor<3x4xf64> {
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64>
%0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [3, 4] : tensor<12xf64> into tensor<3x4xf64>
return %0 : tensor<3x4xf64>
}

func.func @expand_from_sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> {
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64>
%0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [3, 4] : tensor<12xf64, #SparseVector> into tensor<3x4xf64>
return %0 : tensor<3x4xf64>
}

func.func @expand_to_sparse(%arg0: tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> {
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix>
%0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [3, 4] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix>
return %0 : tensor<3x4xf64, #SparseMatrix>
}

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>
%0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [3, 4] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix>
return %0 : tensor<3x4xf64, #SparseMatrix>
}

func.func @expand_dense_3x2x2(%arg0: tensor<3x4xf64>) -> tensor<3x2x2xf64> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<3x4xf64> into tensor<3x2x2xf64>
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [3, 2, 2] : tensor<3x4xf64> into tensor<3x2x2xf64>
return %0 : tensor<3x2x2xf64>
}

func.func @expand_from_sparse_3x2x2(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<3x2x2xf64> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<3x4xf64, #SparseMatrix> into tensor<3x2x2xf64>
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [3, 2, 2] : tensor<3x4xf64, #SparseMatrix> into tensor<3x2x2xf64>
return %0 : tensor<3x2x2xf64>
}

func.func @expand_to_sparse_3x2x2(%arg0: tensor<3x4xf64>) -> tensor<3x2x2xf64, #Sparse3dTensor> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<3x4xf64> into tensor<3x2x2xf64, #Sparse3dTensor>
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [3, 2, 2] : tensor<3x4xf64> into tensor<3x2x2xf64, #Sparse3dTensor>
return %0 : tensor<3x2x2xf64, #Sparse3dTensor>
}

func.func @expand_sparse2sparse_3x2x2(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<3x2x2xf64, #Sparse3dTensor> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<3x4xf64, #SparseMatrix> into tensor<3x2x2xf64, #Sparse3dTensor>
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [3, 2, 2] : tensor<3x4xf64, #SparseMatrix> into tensor<3x2x2xf64, #Sparse3dTensor>
return %0 : tensor<3x2x2xf64, #Sparse3dTensor>
}

func.func @expand_dense_dyn(%arg0: tensor<?x?xf64>) -> tensor<?x2x?xf64> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64> into tensor<?x2x?xf64>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf64>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf64>
%d2 = arith.divui %d1, %c2 : index
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%d0, 2, %d2] : tensor<?x?xf64> into tensor<?x2x?xf64>
return %0 : tensor<?x2x?xf64>
}

func.func @expand_from_sparse_dyn(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64, #SparseMatrix> into tensor<?x2x?xf64>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf64, #SparseMatrix>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf64, #SparseMatrix>
%d2 = arith.divui %d1, %c2 : index
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%d0, 2, %d2] : tensor<?x?xf64, #SparseMatrix> into tensor<?x2x?xf64>
return %0 : tensor<?x2x?xf64>
}

func.func @expand_to_sparse_dyn(%arg0: tensor<?x?xf64>) -> tensor<?x2x?xf64, #Sparse3dTensor> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64> into tensor<?x2x?xf64, #Sparse3dTensor>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf64>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf64>
%d2 = arith.divui %d1, %c2 : index
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%d0, 2, %d2] : tensor<?x?xf64> into tensor<?x2x?xf64, #Sparse3dTensor>
return %0 : tensor<?x2x?xf64, #Sparse3dTensor>
}

func.func @expand_sparse2sparse_dyn(%arg0: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x2x?xf64, #Sparse3dTensor> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] : tensor<?x?xf64, #SparseMatrix> into tensor<?x2x?xf64, #Sparse3dTensor>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf64, #SparseMatrix>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf64, #SparseMatrix>
%d2 = arith.divui %d1, %c2 : index
%0 = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%d0, 2, %d2] : tensor<?x?xf64, #SparseMatrix> into tensor<?x2x?xf64, #Sparse3dTensor>
return %0 : tensor<?x2x?xf64, #Sparse3dTensor>
}

Expand Down

0 comments on commit 7cbaaed

Please sign in to comment.