96a23911f6
A very elaborate, but also very fun revision because all puzzle pieces are finally "falling in place". 1. replaces lingalg annotations + flags with proper sparse tensor types 2. add rigorous verification on sparse tensor type and sparse primitives 3. removes glue and clutter on opaque pointers in favor of sparse tensor types 4. migrates all tests to use sparse tensor types NOTE: next CL will remove *all* obsoleted sparse code in Linalg Reviewed By: bixia Differential Revision: https://reviews.llvm.org/D102095
86 lines
2.6 KiB
MLIR
86 lines
2.6 KiB
MLIR
// RUN: mlir-opt %s \
|
|
// RUN: --sparsification --sparse-tensor-conversion \
|
|
// RUN: --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
|
|
// RUN: --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
|
|
// RUN: --std-bufferize --finalizing-bufferize \
|
|
// RUN: --convert-vector-to-llvm --convert-std-to-llvm | \
|
|
// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
|
|
// RUN: mlir-cpu-runner \
|
|
// RUN: -e entry -entry-point-result=void \
|
|
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
|
|
// RUN: FileCheck %s
|
|
|
|
!Filename = type !llvm.ptr<i8>
|
|
|
|
#SparseMatrix = #sparse_tensor.encoding<{
|
|
dimLevelType = [ "compressed", "compressed" ]
|
|
}>
|
|
|
|
#trait_sum_reduce = {
|
|
indexing_maps = [
|
|
affine_map<(i,j) -> (i,j)>, // A
|
|
affine_map<(i,j) -> ()> // x (out)
|
|
],
|
|
iterator_types = ["reduction", "reduction"],
|
|
doc = "x += A(i,j)"
|
|
}
|
|
|
|
//
|
|
// Integration test that lowers a kernel annotated as sparse to
|
|
// actual sparse code, initializes a matching sparse storage scheme
|
|
// from file, and runs the resulting code with the JIT compiler.
|
|
//
|
|
module {
|
|
//
|
|
// A kernel that sum-reduces a matrix to a single scalar.
|
|
//
|
|
func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix>,
|
|
%argx: tensor<f64>) -> tensor<f64> {
|
|
%0 = linalg.generic #trait_sum_reduce
|
|
ins(%arga: tensor<?x?xf64, #SparseMatrix>)
|
|
outs(%argx: tensor<f64>) {
|
|
^bb(%a: f64, %x: f64):
|
|
%0 = addf %x, %a : f64
|
|
linalg.yield %0 : f64
|
|
} -> tensor<f64>
|
|
return %0 : tensor<f64>
|
|
}
|
|
|
|
func private @getTensorFilename(index) -> (!Filename)
|
|
|
|
//
|
|
// Main driver that reads matrix from file and calls the sparse kernel.
|
|
//
|
|
func @entry() {
|
|
%d0 = constant 0.0 : f64
|
|
%c0 = constant 0 : index
|
|
|
|
// Setup memory for a single reduction scalar,
|
|
// initialized to zero.
|
|
%xdata = memref.alloc() : memref<f64>
|
|
memref.store %d0, %xdata[] : memref<f64>
|
|
%x = memref.tensor_load %xdata : memref<f64>
|
|
|
|
// Read the sparse matrix from file, construct sparse storage.
|
|
%fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
|
|
%a = sparse_tensor.new %fileName : !llvm.ptr<i8> to tensor<?x?xf64, #SparseMatrix>
|
|
|
|
// Call the kernel.
|
|
%0 = call @kernel_sum_reduce(%a, %x)
|
|
: (tensor<?x?xf64, #SparseMatrix>, tensor<f64>) -> tensor<f64>
|
|
|
|
// Print the result for verification.
|
|
//
|
|
// CHECK: 28.2
|
|
//
|
|
%m = memref.buffer_cast %0 : memref<f64>
|
|
%v = memref.load %m[] : memref<f64>
|
|
vector.print %v : f64
|
|
|
|
// Release the resources.
|
|
memref.dealloc %xdata : memref<f64>
|
|
|
|
return
|
|
}
|
|
}
|