forked from OSchip/llvm-project
62 lines
2.2 KiB
MLIR
62 lines
2.2 KiB
MLIR
// RUN: mlir-opt %s -sparsification | \
|
|
// RUN: FileCheck %s --check-prefix=CHECK-SPARSE
|
|
// RUN: mlir-opt %s -sparsification -sparse-tensor-conversion | \
|
|
// RUN: FileCheck %s --check-prefix=CHECK-CONVERT
|
|
|
|
#DCSC = #sparse_tensor.encoding<{
|
|
dimLevelType = [ "compressed", "compressed" ],
|
|
dimOrdering = affine_map<(i,j) -> (j,i)>
|
|
}>
|
|
|
|
#SV = #sparse_tensor.encoding<{
|
|
dimLevelType = [ "compressed" ]
|
|
}>
|
|
|
|
#rowsum = {
|
|
indexing_maps = [
|
|
affine_map<(i,j) -> (i,j)>, // A
|
|
affine_map<(i,j) -> (i)> // x (out)
|
|
],
|
|
iterator_types = ["parallel", "reduction"],
|
|
doc = "X(i) = SUM A(i,j)"
|
|
}
|
|
|
|
//
|
|
// CHECK-SPARSE-LABEL: func @kernel(
|
|
// CHECK-SPARSE: %[[A:.*]], %[[B:.*]], %[[C:.*]], %{{.*}} = sparse_tensor.expand
|
|
// CHECK-SPARSE: scf.for
|
|
// CHECK-SPARSE: scf.for
|
|
// CHECK-SPARSE: sparse_tensor.compress %{{.*}}, %{{.*}}, %[[A]], %[[B]], %[[C]]
|
|
// CHECK-SPARSE: %[[RET:.*]] = sparse_tensor.load %{{.*}} hasInserts
|
|
// CHECK-SPARSE: return %[[RET]]
|
|
//
|
|
// CHECK-CONVERT-LABEL: func @kernel(
|
|
// CHECK-CONVERT: %{{.*}} = call @sparseDimSize
|
|
// CHECK-CONVERT: %[[S:.*]] = call @sparseDimSize
|
|
// CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
|
|
// CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
|
|
// CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
|
|
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>)
|
|
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
|
|
// CHECK-CONVERT: scf.for
|
|
// CHECK-CONVERT: scf.for
|
|
// CHECK-CONVERT: call @expInsertF64
|
|
// CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64>
|
|
// CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1>
|
|
// CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex>
|
|
// CHECK-CONVERT: call @endInsert
|
|
//
|
|
func.func @kernel(%arga: tensor<?x?xf64, #DCSC>) -> tensor<?xf64, #SV> {
|
|
%c0 = arith.constant 0 : index
|
|
%n = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSC>
|
|
%v = sparse_tensor.init [%n] : tensor<?xf64, #SV>
|
|
%0 = linalg.generic #rowsum
|
|
ins(%arga: tensor<?x?xf64, #DCSC>)
|
|
outs(%v: tensor<?xf64, #SV>) {
|
|
^bb(%a: f64, %x: f64):
|
|
%1 = arith.addf %x, %a : f64
|
|
linalg.yield %1 : f64
|
|
} -> tensor<?xf64, #SV>
|
|
return %0 : tensor<?xf64, #SV>
|
|
}
|