[Mlir-commits] [mlir] 3420cd7 - [mlir][sparse] Add testing for bf16 and fallback for software bf16
Benjamin Kramer
llvmlistbot at llvm.org
Fri Jun 17 12:54:14 PDT 2022
Author: Benjamin Kramer
Date: 2022-06-17T21:54:01+02:00
New Revision: 3420cd7caf8d29e2691bb4a6769425bb033a1224
URL: https://github.com/llvm/llvm-project/commit/3420cd7caf8d29e2691bb4a6769425bb033a1224
DIFF: https://github.com/llvm/llvm-project/commit/3420cd7caf8d29e2691bb4a6769425bb033a1224.diff
LOG: [mlir][sparse] Add testing for bf16 and fallback for software bf16
This adds weak versions of the truncation libcalls in case the runtime
environment doesn't have them.
Differential Revision: https://reviews.llvm.org/D128091
Added:
mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir
Modified:
mlir/lib/ExecutionEngine/Float16bits.cpp
Removed:
################################################################################
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<float>(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<?xbf16, #SparseVector>,
+ %argb: tensor<?xbf16, #SparseVector>) -> tensor<?xbf16, #DenseVector> {
+ %c = arith.constant 0 : index
+ %d = tensor.dim %arga, %c : tensor<?xbf16, #SparseVector>
+ %xv = bufferization.alloc_tensor (%d) : tensor<?xbf16, #DenseVector>
+ %0 = linalg.generic #trait_vec_op
+ ins(%arga, %argb: tensor<?xbf16, #SparseVector>, tensor<?xbf16, #SparseVector>)
+ outs(%xv: tensor<?xbf16, #DenseVector>) {
+ ^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<?xbf16, #DenseVector>
+ return %0 : tensor<?xbf16, #DenseVector>
+ }
+
+ // Dumps a dense vector of type bf16.
+ func.func @dump_vec(%arg0: tensor<?xbf16, #DenseVector>) {
+ // 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<?xbf16, #DenseVector> to memref<?xbf16>
+ %1 = vector.transfer_read %0[%c0], %d0: memref<?xbf16>, 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<?xbf16, #SparseVector>
+ %sv2 = sparse_tensor.convert %v2 : tensor<32xbf16> to tensor<?xbf16, #SparseVector>
+
+ // Call the sparse vector kernel.
+ %0 = call @vector_min(%sv1, %sv2)
+ : (tensor<?xbf16, #SparseVector>,
+ tensor<?xbf16, #SparseVector>) -> tensor<?xbf16, #DenseVector>
+
+ //
+ // 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<?xbf16, #DenseVector>) -> ()
+
+ // Release the resources.
+ sparse_tensor.release %sv1 : tensor<?xbf16, #SparseVector>
+ sparse_tensor.release %sv2 : tensor<?xbf16, #SparseVector>
+ sparse_tensor.release %0 : tensor<?xbf16, #DenseVector>
+ 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<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)"
+}
+
+module {
+ //
+ // A kernel that sum-reduces a matrix to a single scalar.
+ //
+ func.func @kernel_sum_reduce(%arga: tensor<?x?xbf16, #SparseMatrix>,
+ %argx: tensor<bf16> {linalg.inplaceable = true}) -> tensor<bf16> {
+ %0 = linalg.generic #trait_sum_reduce
+ ins(%arga: tensor<?x?xbf16, #SparseMatrix>)
+ outs(%argx: tensor<bf16>) {
+ ^bb(%a: bf16, %x: bf16):
+ %0 = arith.addf %x, %a : bf16
+ linalg.yield %0 : bf16
+ } -> tensor<bf16>
+ return %0 : tensor<bf16>
+ }
+
+ 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<?x?xbf16, #SparseMatrix>
+
+ %d0 = arith.constant 0.0 : bf16
+ // Setup memory for a single reduction scalar,
+ // initialized to zero.
+ %xdata = memref.alloc() : memref<bf16>
+ memref.store %d0, %xdata[] : memref<bf16>
+ %x = bufferization.to_tensor %xdata : memref<bf16>
+
+ // Call the kernel.
+ %0 = call @kernel_sum_reduce(%a, %x)
+ : (tensor<?x?xbf16, #SparseMatrix>, tensor<bf16>) -> tensor<bf16>
+
+ // Print the result for verification.
+ //
+ // CHECK: 13.5
+ //
+ %m = bufferization.to_memref %0 : memref<bf16>
+ %v = memref.load %m[] : memref<bf16>
+ %vf = arith.extf %v: bf16 to f32
+ vector.print %vf : f32
+
+ // Release the resources.
+ memref.dealloc %xdata : memref<bf16>
+ sparse_tensor.release %a : tensor<?x?xbf16, #SparseMatrix>
+
+ return
+ }
+}
More information about the Mlir-commits
mailing list