//--------------------------------------------------------------------------------------------------
// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
//
// Set-up that's shared across all tests in this directory. In principle, this
// config could be moved to lit.local.cfg. However, there are downstream users that
// do not use these LIT config files. Hence why this is kept inline.
//
// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}
//
// DEFINE: %{env} =
//--------------------------------------------------------------------------------------------------
// RUN: %{compile} | %{run} | FileCheck %s
//
// Do the same run, but now with direct IR generation.
// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true
// RUN: %{compile} | %{run} | FileCheck %s
//
// Do the same run, but now with direct IR generation and vectorization.
// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true
// RUN: %{compile} | %{run} | FileCheck %s
#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
#DV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : dense) }>
#trait_reduction = {
indexing_maps = [
affine_map<(i) -> (i)>, // a
affine_map<(i) -> ()> // x (scalar out)
],
iterator_types = ["reduction"],
doc = "x += PROD_CUSTOM_i a(i)"
}
// An example of vector reductions.
module {
// Custom prod reduction: stored i32 elements only.
func.func @prod_dreduction_i32(%arga: tensor<32xi32, #DV>,
%argx: tensor<i32>) -> tensor<i32> {
%c = tensor.extract %argx[] : tensor<i32>
%0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #DV>)
outs(%argx: tensor<i32>) {
^bb(%a: i32, %b: i32):
%1 = sparse_tensor.reduce %a, %b, %c : i32 {
^bb0(%x: i32, %y: i32):
%2 = arith.muli %x, %y : i32
sparse_tensor.yield %2 : i32
}
linalg.yield %1 : i32
} -> tensor<i32>
return %0 : tensor<i32>
}
// Custom prod reduction: stored f32 elements only.
func.func @prod_dreduction_f32(%arga: tensor<32xf32, #DV>,
%argx: tensor<f32>) -> tensor<f32> {
%c = tensor.extract %argx[] : tensor<f32>
%0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xf32, #DV>)
outs(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32):
%1 = sparse_tensor.reduce %a, %b, %c : f32 {
^bb0(%x: f32, %y: f32):
%2 = arith.mulf %x, %y : f32
sparse_tensor.yield %2 : f32
}
linalg.yield %1 : f32
} -> tensor<f32>
return %0 : tensor<f32>
}
// Custom prod reduction: stored i32 elements only.
func.func @prod_sreduction_i32(%arga: tensor<32xi32, #SV>,
%argx: tensor<i32>) -> tensor<i32> {
%c = tensor.extract %argx[] : tensor<i32>
%0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #SV>)
outs(%argx: tensor<i32>) {
^bb(%a: i32, %b: i32):
%1 = sparse_tensor.reduce %a, %b, %c : i32 {
^bb0(%x: i32, %y: i32):
%2 = arith.muli %x, %y : i32
sparse_tensor.yield %2 : i32
}
linalg.yield %1 : i32
} -> tensor<i32>
return %0 : tensor<i32>
}
// Custom prod reduction: stored f32 elements only.
func.func @prod_sreduction_f32(%arga: tensor<32xf32, #SV>,
%argx: tensor<f32>) -> tensor<f32> {
%c = tensor.extract %argx[] : tensor<f32>
%0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xf32, #SV>)
outs(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32):
%1 = sparse_tensor.reduce %a, %b, %c : f32 {
^bb0(%x: f32, %y: f32):
%2 = arith.mulf %x, %y : f32
sparse_tensor.yield %2 : f32
}
linalg.yield %1 : f32
} -> tensor<f32>
return %0 : tensor<f32>
}
// Custom prod reduction: stored i32 elements and implicit zeros.
//
// NOTE: this is a somewhat strange operation, since for most sparse
// situations the outcome would always be zero; it is added
// to test full functionality and illustrate the subtle differences
// between the various custom operations; it would make a bit more
// sense for e.g. a min/max reductions, although it still would
// "densify" the iteration space.
//
func.func @prod_xreduction_i32(%arga: tensor<32xi32, #SV>,
%argx: tensor<i32>) -> tensor<i32> {
%c = tensor.extract %argx[] : tensor<i32>
%0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #SV>)
outs(%argx: tensor<i32>) {
^bb(%a: i32, %b: i32):
%u = sparse_tensor.unary %a : i32 to i32
present={
^bb0(%x: i32):
sparse_tensor.yield %x : i32
} absent={
^bb0:
%c0 = arith.constant 0 : i32
sparse_tensor.yield %c0 : i32
}
%1 = sparse_tensor.reduce %u, %b, %c : i32 {
^bb0(%x: i32, %y: i32):
%2 = arith.muli %x, %y : i32
sparse_tensor.yield %2 : i32
}
linalg.yield %1 : i32
} -> tensor<i32>
return %0 : tensor<i32>
}
func.func @dump_i32(%arg0 : tensor<i32>) {
%v = tensor.extract %arg0[] : tensor<i32>
vector.print %v : i32
return
}
func.func @dump_f32(%arg0 : tensor<f32>) {
%v = tensor.extract %arg0[] : tensor<f32>
vector.print %v : f32
return
}
func.func @main() {
// Note: Constants bufferize to read-only buffers.
%ri = arith.constant dense< 7 > : tensor<i32>
%rf = arith.constant dense< 2.0 > : tensor<f32>
// Vectors with a few zeros.
%c_0_i32 = arith.constant dense<[
1, 1, 7, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 3, 0, 1, 1, 1, 1, 1, 0, 1, 1, 7, 3
]> : tensor<32xi32>
%c_0_f32 = arith.constant dense<[
1.0, 1.0, 1.0, 3.5, 1.0, 1.0, 1.0, 1.0,
1.0, 0.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0,
1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0,
1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0
]> : tensor<32xf32>
// Vectors with no zeros.
%c_1_i32 = arith.constant dense<[
1, 1, 7, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 3
]> : tensor<32xi32>
%c_1_f32 = arith.constant dense<[
1.0, 1.0, 1.0, 3.5, 1.0, 1.0, 1.0, 1.0,
1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0,
1.0, 1.0, 1.0, 1.0, 3.0, 1.0, 1.0, 1.0,
1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 4.0
]> : tensor<32xf32>
// Convert constants to annotated tensors. Note that this
// particular conversion only stores nonzero elements,
// so we will have no explicit zeros, only implicit zeros.
%d0_i32 = sparse_tensor.convert %c_0_i32
: tensor<32xi32> to tensor<32xi32, #DV>
%d0_f32 = sparse_tensor.convert %c_0_f32
: tensor<32xf32> to tensor<32xf32, #DV>
%s0_i32 = sparse_tensor.convert %c_0_i32
: tensor<32xi32> to tensor<32xi32, #SV>
%s0_f32 = sparse_tensor.convert %c_0_f32
: tensor<32xf32> to tensor<32xf32, #SV>
%d1_i32 = sparse_tensor.convert %c_1_i32
: tensor<32xi32> to tensor<32xi32, #DV>
%d1_f32 = sparse_tensor.convert %c_1_f32
: tensor<32xf32> to tensor<32xf32, #DV>
%s1_i32 = sparse_tensor.convert %c_1_i32
: tensor<32xi32> to tensor<32xi32, #SV>
%s1_f32 = sparse_tensor.convert %c_1_f32
: tensor<32xf32> to tensor<32xf32, #SV>
// Special case, construct a sparse vector with an explicit zero.
%v0 = arith.constant sparse< [ [1] ], [ 0 ] > : tensor<32xi32>
%s0 = sparse_tensor.convert %v0: tensor<32xi32> to tensor<32xi32, #SV>
// Call the kernels.
%0 = call @prod_dreduction_i32(%d0_i32, %ri) : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
%1 = call @prod_dreduction_f32(%d0_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
%2 = call @prod_sreduction_i32(%s0_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
%3 = call @prod_sreduction_f32(%s0_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
%4 = call @prod_dreduction_i32(%d1_i32, %ri) : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
%5 = call @prod_dreduction_f32(%d1_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
%6 = call @prod_sreduction_i32(%s1_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
%7 = call @prod_sreduction_f32(%s1_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
%8 = call @prod_sreduction_i32(%s0, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
%9 = call @prod_xreduction_i32(%s0_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
%10 = call @prod_xreduction_i32(%s1_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
// Verify results. Note that the custom reduction gave permission
// to treat an explicit vs implicit zero differently to compute the
// full product reduction over stored elements. A "standard" product
// reduction would have to return 0 for any implicit zero occurrence
// too. An explicit zero nullifies the product, though, as requested.
//
// CHECK: 0
// CHECK: 0
// CHECK: 3087
// CHECK: 14
// CHECK: 3087
// CHECK: 168
// CHECK: 3087
// CHECK: 168
// CHECK: 0
// CHECK: 0
// CHECK: 3087
//
call @dump_i32(%0) : (tensor<i32>) -> ()
call @dump_f32(%1) : (tensor<f32>) -> ()
call @dump_i32(%2) : (tensor<i32>) -> ()
call @dump_f32(%3) : (tensor<f32>) -> ()
call @dump_i32(%4) : (tensor<i32>) -> ()
call @dump_f32(%5) : (tensor<f32>) -> ()
call @dump_i32(%6) : (tensor<i32>) -> ()
call @dump_f32(%7) : (tensor<f32>) -> ()
call @dump_i32(%8) : (tensor<i32>) -> ()
call @dump_i32(%9) : (tensor<i32>) -> ()
call @dump_i32(%10) : (tensor<i32>) -> ()
// Release the resources.
bufferization.dealloc_tensor %d0_i32 : tensor<32xi32, #DV>
bufferization.dealloc_tensor %d0_f32 : tensor<32xf32, #DV>
bufferization.dealloc_tensor %s0_i32 : tensor<32xi32, #SV>
bufferization.dealloc_tensor %s0_f32 : tensor<32xf32, #SV>
bufferization.dealloc_tensor %d1_i32 : tensor<32xi32, #DV>
bufferization.dealloc_tensor %d1_f32 : tensor<32xf32, #DV>
bufferization.dealloc_tensor %s1_i32 : tensor<32xi32, #SV>
bufferization.dealloc_tensor %s1_f32 : tensor<32xf32, #SV>
bufferization.dealloc_tensor %s0 : tensor<32xi32, #SV>
bufferization.dealloc_tensor %0 : tensor<i32>
bufferization.dealloc_tensor %1 : tensor<f32>
bufferization.dealloc_tensor %2 : tensor<i32>
bufferization.dealloc_tensor %3 : tensor<f32>
bufferization.dealloc_tensor %4 : tensor<i32>
bufferization.dealloc_tensor %5 : tensor<f32>
bufferization.dealloc_tensor %6 : tensor<i32>
bufferization.dealloc_tensor %7 : tensor<f32>
bufferization.dealloc_tensor %8 : tensor<i32>
bufferization.dealloc_tensor %9 : tensor<i32>
bufferization.dealloc_tensor %10 : tensor<i32>
return
}
}
Codebase Browser
llvm