[mlir][sparse] Add testing for bf16 and fallback for software bf16

[bkramer]

Summary

This adds weak versions of the truncation libcalls in case the runtime
environment doesn't have them.

[aartbik]

mlir/lib/ExecutionEngine/Float16bits.cpp

146

here and below, just curious if you meant "routine", or really "routing" (as in, different path)

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir

2

left indent

mlir/lib/ExecutionEngine/Float16bits.cpp

 //===--- Float16bits.cpp - supports 2-byte floats  ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements f16 and bf16 to support the compilation and execution
 // of programs using these types.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/Float16bits.h"
+#include "llvm/Support/Compiler.h"
 
 namespace {
 
 // Union used to make the int/float aliasing explicit so we can access the raw
 // bits.
 union Float32Bits {
   uint32_t u;
   float f;
std::ostream &operator<<(std::ostream &os, const f16 &f) {
   os << half2float(f.bits);
   return os;
 }
 
 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));
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir

+// 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
+  }
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir

+// 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
+  }
+}