diff --git a/mlir/lib/ExecutionEngine/Float16bits.cpp b/mlir/lib/ExecutionEngine/Float16bits.cpp index f9f2bbc7333c8..c376022c69dca 100644 --- a/mlir/lib/ExecutionEngine/Float16bits.cpp +++ b/mlir/lib/ExecutionEngine/Float16bits.cpp @@ -12,6 +12,7 @@ //===----------------------------------------------------------------------===// #include "mlir/ExecutionEngine/Float16bits.h" +#include "llvm/Support/Compiler.h" namespace { @@ -141,3 +142,17 @@ std::ostream &operator<<(std::ostream &os, const bf16 &d) { os << bfloat2float(d.bits); return os; } + +// Provide a float->bfloat conversion routine in case the runtime doesn't have +// one. +extern "C" uint16_t LLVM_ATTRIBUTE_WEAK __truncsfbf2(float f) { + return float2bfloat(f); +} + +// Provide a double->bfloat conversion routine in case the runtime doesn't have +// one. +extern "C" uint16_t LLVM_ATTRIBUTE_WEAK __truncdfbf2(double d) { + // This does a double rounding step, but it's precise enough for our use + // cases. + return __truncsfbf2(static_cast(d)); +} diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir new file mode 100644 index 0000000000000..f2c595eb35d9a --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir @@ -0,0 +1,90 @@ +// RUN: mlir-opt %s --sparse-compiler | \ +// 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 + +#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}> +#DenseVector = #sparse_tensor.encoding<{dimLevelType = ["dense"]}> + +#trait_vec_op = { + indexing_maps = [ + affine_map<(i) -> (i)>, // a (in) + affine_map<(i) -> (i)>, // b (in) + affine_map<(i) -> (i)> // x (out) + ], + iterator_types = ["parallel"] +} + +module { + // Creates a dense vector using the minimum values from two input sparse vectors. + // When there is no overlap, include the present value in the output. + func.func @vector_min(%arga: tensor, + %argb: tensor) -> tensor { + %c = arith.constant 0 : index + %d = tensor.dim %arga, %c : tensor + %xv = bufferization.alloc_tensor (%d) : tensor + %0 = linalg.generic #trait_vec_op + ins(%arga, %argb: tensor, tensor) + outs(%xv: tensor) { + ^bb(%a: bf16, %b: bf16, %x: bf16): + %1 = sparse_tensor.binary %a, %b : bf16, bf16 to bf16 + overlap={ + ^bb0(%a0: bf16, %b0: bf16): + %cmp = arith.cmpf "olt", %a0, %b0 : bf16 + %2 = arith.select %cmp, %a0, %b0: bf16 + sparse_tensor.yield %2 : bf16 + } + left=identity + right=identity + linalg.yield %1 : bf16 + } -> tensor + return %0 : tensor + } + + // Dumps a dense vector of type bf16. + func.func @dump_vec(%arg0: tensor) { + // Dump the values array to verify only sparse contents are stored. + %c0 = arith.constant 0 : index + %d0 = arith.constant -1.0 : bf16 + %0 = sparse_tensor.values %arg0 : tensor to memref + %1 = vector.transfer_read %0[%c0], %d0: memref, vector<32xbf16> + %f1 = arith.extf %1: vector<32xbf16> to vector<32xf32> + vector.print %f1 : vector<32xf32> + return + } + + // Driver method to call and verify the kernel. + func.func @entry() { + %c0 = arith.constant 0 : index + + // Setup sparse vectors. + %v1 = arith.constant sparse< + [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ], + [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] + > : tensor<32xbf16> + %v2 = arith.constant sparse< + [ [1], [3], [4], [10], [16], [18], [21], [28], [29], [31] ], + [11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0 ] + > : tensor<32xbf16> + %sv1 = sparse_tensor.convert %v1 : tensor<32xbf16> to tensor + %sv2 = sparse_tensor.convert %v2 : tensor<32xbf16> to tensor + + // Call the sparse vector kernel. + %0 = call @vector_min(%sv1, %sv2) + : (tensor, + tensor) -> tensor + + // + // Verify the result. + // + // CHECK: ( 1, 11, 0, 2, 13, 0, 0, 0, 0, 0, 14, 3, 0, 0, 0, 0, 15, 4, 16, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9 ) + call @dump_vec(%0) : (tensor) -> () + + // Release the resources. + sparse_tensor.release %sv1 : tensor + sparse_tensor.release %sv2 : tensor + sparse_tensor.release %0 : tensor + return + } +} diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir new file mode 100644 index 0000000000000..6e9f8b17eff84 --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir @@ -0,0 +1,78 @@ +// RUN: mlir-opt %s --sparse-compiler | \ +// 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 = !llvm.ptr + +#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)" +} + +module { + // + // A kernel that sum-reduces a matrix to a single scalar. + // + func.func @kernel_sum_reduce(%arga: tensor, + %argx: tensor {linalg.inplaceable = true}) -> tensor { + %0 = linalg.generic #trait_sum_reduce + ins(%arga: tensor) + outs(%argx: tensor) { + ^bb(%a: bf16, %x: bf16): + %0 = arith.addf %x, %a : bf16 + linalg.yield %0 : bf16 + } -> tensor + return %0 : tensor + } + + func.func private @getTensorFilename(index) -> (!Filename) + + // + // Main driver that reads matrix from file and calls the sparse kernel. + // + func.func @entry() { + // Setup input sparse matrix from compressed constant. + %d = arith.constant dense <[ + [ 1.1, 1.2, 0.0, 1.4 ], + [ 0.0, 0.0, 0.0, 0.0 ], + [ 3.1, 0.0, 3.3, 3.4 ] + ]> : tensor<3x4xbf16> + %a = sparse_tensor.convert %d : tensor<3x4xbf16> to tensor + + %d0 = arith.constant 0.0 : bf16 + // Setup memory for a single reduction scalar, + // initialized to zero. + %xdata = memref.alloc() : memref + memref.store %d0, %xdata[] : memref + %x = bufferization.to_tensor %xdata : memref + + // Call the kernel. + %0 = call @kernel_sum_reduce(%a, %x) + : (tensor, tensor) -> tensor + + // Print the result for verification. + // + // CHECK: 13.5 + // + %m = bufferization.to_memref %0 : memref + %v = memref.load %m[] : memref + %vf = arith.extf %v: bf16 to f32 + vector.print %vf : f32 + + // Release the resources. + memref.dealloc %xdata : memref + sparse_tensor.release %a : tensor + + return + } +}