[mlir][sparse] Add F16 and BF16.

[bixia]

Summary

This is the first PR to add F16 and BF16 support to the sparse codegen. There are still problems in supporting these two data types, such as BF16 is not quite working yet.

Add tests cases.

[bixia]

Add CMakeLists.txt and BUILD.bazel change.

[aartbik]

mlir/include/mlir/ExecutionEngine/Float16bits.h

2

nit: typically, more of the empty white space is filled with the ------- line here

11

them -> these types

21–22

The reference to the FOREVERY macro seems a bit out of in place in this file (given that others may use it outside the sparse compiler context). I would omit it, or, if you want to keep it, otherwise put it in a separate sentence with some intro explanation.

35

Add documentation on providing the output methods for both.

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

81–88

can we keep FP before INT, just like we do at all other places?

mlir/lib/ExecutionEngine/Float16bits.cpp

2

same nit: fill the whitespace with -----------------------===//

11

them -> these types

19

Add documentation, i.e. union used to make the int/float aliasing explicit so we can access the raw bits

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1570

Can we say,

// Two-byte floats

just so we start with a capital letter

[bixia]

Address review comments.

mlir/include/mlir/ExecutionEngine/Float16bits.h

21–22

Removed the reference.

[wrengr]

mlir/lib/ExecutionEngine/Float16bits.cpp

36–37

Why not combine these into f.u ^= (f.u & signMask); ?

72–73

ditto earlier comment about using constexpr

76

You can (and should) drop the equals sign here, to ensure there's no extra temporary value constructed

78–83

I feel like these (112 <<23) and (113 << 23) expression should be factored out into constexpr variables to be shared throughout the file. Especially since doing so means giving them names to indicate what they mean/do.

119

why not f16::f16(float f) : bits(float2half(f)) {} ?

121

ditto

[bixia]

Fixed a test.

[aartbik]

mlir/lib/ExecutionEngine/Float16bits.cpp

78–83

+1 on that. Also, probably the shift constants (13, 23, 16) could get some symbolic names

[bixia]

Replaced some constants with names.

[aartbik]

I think something went wrong. In one revision, you had addressed comments but in a later upload, some of these were undone again?

[bixia]

Restore changes that addressed the first round of comments.

mlir/lib/ExecutionEngine/Float16bits.cpp

36–37

sign is used in line 64

72–73

I think you are referring to the comment for line 76, will wait for your response there to clarify the suggestion.

76

Sorry, I don't understand the suggestion here. Would you please spell out the whole statement for this line here?

78–83

Added kF32Magic, kF32HalfExpAdjust and a few others.

[aartbik]

One last nit. Also, wait a bit before submitting to see if Wren has some more feedback on addressing the comment she made earlier.

mlir/lib/ExecutionEngine/Float16bits.cpp

48

if you give these two hex numbers better symbolic names, it reads even better

[bixia]

Name the constants for qNan and Inf.

[wrengr]

Sorry I missed the comments; I was out yesterday for a doctor's appt.

I'll make a followup CL with the changes I had in mind. While working on that I recall that we might not be able to use some of my suggestions, since they're C++11 features whereas this file must be C++98 compliant. Nevertheless, once I finish the CL I'll post it just to help clarify what I meant.

[mehdi_amini]

mlir/lib/ExecutionEngine/Float16bits.cpp

116

Isn't this kind of type punning actually UB?

[mehdi_amini]

mlir/lib/ExecutionEngine/Float16bits.cpp

32

I can't find a reference in Eigen, I found https://gitlab.com/libeigen/eigen/-/blob/master/Eigen/src/Core/arch/Default/BFloat16.h#L347-359 which is much simpler actually. Can you help?

mlir/include/mlir/ExecutionEngine/Float16bits.h

+//===--- Float16bits.h - 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.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_EXECUTIONENGINE_FLOAT16BITS_H_
+#define MLIR_EXECUTIONENGINE_FLOAT16BITS_H_
+
+#include <cstdint>
+#include <iostream>
+
+// Implements half precision and bfloat with f16 and bf16, using the MLIR type
+// names. These data types are also used for c-interface runtime routines.
+extern "C" {
+struct f16 {
+  f16(float f = 0);
+  uint16_t bits;
+};
+
+struct bf16 {
+  bf16(float f = 0);
+  uint16_t bits;
+};
+}
+
+// Outputs a half precision value.
+std::ostream &operator<<(std::ostream &os, const f16 &f);
+// Outputs a bfloat value.
+std::ostream &operator<<(std::ostream &os, const bf16 &d);
+
+#endif // MLIR_EXECUTIONENGINE_FLOAT16BITS_H_

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 // public API of `ExecutionEngine/SparseUtils.cpp`.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
 #define MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
+#include "mlir/ExecutionEngine/Float16bits.h"
 
 #include <cinttypes>
 #include <complex>
 #include <vector>
 
 extern "C" {
 
 //===----------------------------------------------------------------------===//

 // or MLIR's `ComplexType`).
 using complex64 = std::complex<double>;
 using complex32 = std::complex<float>;
 
 /// Encoding of the elemental type, for "overloading" @newSparseTensor.
 enum class PrimaryType : uint32_t {
   kF64 = 1,
   kF32 = 2,
-  kI64 = 3,
-  kI32 = 4,
-  kI16 = 5,
-  kI8 = 6,
-  kC64 = 7,
-  kC32 = 8
+  kF16 = 3,
+  kBF16 = 4,
+  kI64 = 5,
+  kI32 = 6,
+  kI16 = 7,
+  kI8 = 8,
+  kC64 = 9,
+  kC32 = 10
 };
 
 // This x-macro only specifies the non-complex `V` types, because the ABI
 // for complex types has compiler-/architecture-dependent details we need
 // to work around.  Namely, when a function takes a parameter of C/C++
 // type `complex32` (per se), then there is additional padding that causes
 // it not to match the LLVM type `!llvm.struct<(f32, f32)>`.  This only
 // happens with the `complex32` type itself, not with pointers/arrays
 // of complex values.  We also exclude `complex64` because it is in
 // principle susceptible to analogous ABI issues (even though we haven't
 // yet encountered them in practice).
 #define FOREVERY_SIMPLEX_V(DO)                                                 \
   DO(F64, double)                                                              \
   DO(F32, float)                                                               \
+  DO(F16, f16)                                                                 \
+  DO(BF16, bf16)                                                               \
   DO(I64, int64_t)                                                             \
   DO(I32, int32_t)                                                             \
   DO(I16, int16_t)                                                             \
   DO(I8, int8_t)
 
 // This x-macro includes all `V` types, for when the aforementioned ABI
 // issues don't apply (e.g., because the functions take pointers/arrays).
 #define FOREVERY_V(DO)                                                         \

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(Type tp) {
   return overheadTypeFunctionSuffix(overheadTypeEncoding(tp));
 }
 
 PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
   if (elemTp.isF64())
     return PrimaryType::kF64;
   if (elemTp.isF32())
     return PrimaryType::kF32;
+  if (elemTp.isF16())
+    return PrimaryType::kF16;
+  if (elemTp.isBF16())
+    return PrimaryType::kBF16;
   if (elemTp.isInteger(64))
     return PrimaryType::kI64;
   if (elemTp.isInteger(32))
     return PrimaryType::kI32;
   if (elemTp.isInteger(16))
     return PrimaryType::kI16;
   if (elemTp.isInteger(8))
     return PrimaryType::kI8;

mlir/lib/ExecutionEngine/CMakeLists.txt

 # Exclude these from libMLIR.so because the JIT infrastructure
 # is a big dependency which most don't need.
 
 set(LLVM_OPTIONAL_SOURCES
   AsyncRuntime.cpp
   CRunnerUtils.cpp
   CudaRuntimeWrappers.cpp
   SparseTensorUtils.cpp
   ExecutionEngine.cpp
+  Float16bits.cpp
   RocmRuntimeWrappers.cpp
   RunnerUtils.cpp
   OptUtils.cpp
   JitRunner.cpp
   )
 
 # Use a separate library for OptUtils, to avoid pulling in the entire JIT and
 # codegen infrastructure. Unlike MLIRExecutionEngine, this is part of
add_mlir_library(MLIRJitRunner
   MLIRTargetLLVMIRExport
   MLIRTransforms
   MLIRSupport
 )
 
 add_mlir_library(mlir_c_runner_utils
   SHARED
   CRunnerUtils.cpp
+  Float16bits.cpp
   SparseTensorUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
   )
 set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11)
 target_compile_definitions(mlir_c_runner_utils PRIVATE mlir_c_runner_utils_EXPORTS)
 
 add_mlir_library(mlir_runner_utils

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"
+
+namespace {
+
+// Union used to make the int/float aliasing explicit so we can access the raw
+// bits.
+union Float32Bits {
+  uint32_t u;
+  float f;
+};
+
+const uint32_t kF32MantiBits = 23;
+const uint32_t kF32HalfMantiBitDiff = 13;
+const uint32_t kF32HalfBitDiff = 16;
+const Float32Bits kF32Magic = {113 << kF32MantiBits};
+const uint32_t kF32HalfExpAdjust = (127 - 15) << kF32MantiBits;
+
+// Constructs the 16 bit representation for a half precision value from a float
+// value. This implementation is adapted from Eigen.
+uint16_t float2half(float floatValue) {
+  const Float32Bits inf = {255 << kF32MantiBits};
+  const Float32Bits f16max = {(127 + 16) << kF32MantiBits};
+  const Float32Bits denormMagic = {((127 - 15) + (kF32MantiBits - 10) + 1)
+                                   << kF32MantiBits};
+  uint32_t signMask = 0x80000000u;
+  uint16_t halfValue = static_cast<uint16_t>(0x0u);
+  Float32Bits f;
+  f.f = floatValue;
+  uint32_t sign = f.u & signMask;
+  f.u ^= sign;
+
+  if (f.u >= f16max.u) {
+    const uint32_t halfQnan = 0x7e00;
+    const uint32_t halfInf = 0x7c00;
+    // Inf or NaN (all exponent bits set).
+    halfValue = (f.u > inf.u) ? halfQnan : halfInf; // NaN->qNaN and Inf->Inf
+  } else {
+    // (De)normalized number or zero.
+    if (f.u < kF32Magic.u) {
+      // The resulting FP16 is subnormal or zero.
+      //
+      // Use a magic value to align our 10 mantissa bits at the bottom of the
+      // float. As long as FP addition is round-to-nearest-even this works.
+      f.f += denormMagic.f;
+
+      halfValue = static_cast<uint16_t>(f.u - denormMagic.u);
+    } else {
+      uint32_t mantOdd =
+          (f.u >> kF32HalfMantiBitDiff) & 1; // Resulting mantissa is odd.
+
+      // Update exponent, rounding bias part 1. The following expressions are
+      // equivalent to `f.u += ((unsigned int)(15 - 127) << kF32MantiBits) +
+      // 0xfff`, but without arithmetic overflow.
+      f.u += 0xc8000fffU;
+      // Rounding bias part 2.
+      f.u += mantOdd;
+      halfValue = static_cast<uint16_t>(f.u >> kF32HalfMantiBitDiff);
+    }
+  }
+
+  halfValue |= static_cast<uint16_t>(sign >> kF32HalfBitDiff);
+  return halfValue;
+}
+
+// Converts the 16 bit representation of a half precision value to a float
+// value. This implementation is adapted from Eigen.
+float half2float(uint16_t halfValue) {
+  const uint32_t shiftedExp =
+      0x7c00 << kF32HalfMantiBitDiff; // Exponent mask after shift.
+
+  // Initialize the float representation with the exponent/mantissa bits.
+  Float32Bits f = {
+      static_cast<uint32_t>((halfValue & 0x7fff) << kF32HalfMantiBitDiff)};
+  const uint32_t exp = shiftedExp & f.u;
+  f.u += kF32HalfExpAdjust; // Adjust the exponent
+
+  // Handle exponent special cases.
+  if (exp == shiftedExp) {
+    // Inf/NaN
+    f.u += kF32HalfExpAdjust;
+  } else if (exp == 0) {
+    // Zero/Denormal?
+    f.u += 1 << kF32MantiBits;
+    f.f -= kF32Magic.f;
+  }
+
+  f.u |= (halfValue & 0x8000) << kF32HalfBitDiff; // Sign bit.
+  return f.f;
+}
+
+const uint32_t kF32BfMantiBitDiff = 16;
+
+// Constructs the 16 bit representation for a bfloat value from a float value.
+// This implementation is adapted from Eigen.
+uint16_t float2bfloat(float floatValue) {
+  Float32Bits floatBits;
+  floatBits.f = floatValue;
+  uint16_t bfloatBits;
+
+  // Least significant bit of resulting bfloat.
+  uint32_t lsb = (floatBits.u >> kF32BfMantiBitDiff) & 1;
+  uint32_t rounding_bias = 0x7fff + lsb;
+  floatBits.u += rounding_bias;
+  bfloatBits = static_cast<uint16_t>(floatBits.u >> kF32BfMantiBitDiff);
+  return bfloatBits;
+}
+
+// Converts the 16 bit representation of a bfloat value to a float value. This
+// implementation is adapted from Eigen.
+float bfloat2float(uint16_t bfloatBits) {
+  Float32Bits floatBits;
+  floatBits.u = static_cast<uint32_t>(bfloatBits) << kF32BfMantiBitDiff;
+  return floatBits.f;
+}
+
+} // namespace
+
+f16::f16(float f) : bits(float2half(f)) {}
+
+bf16::bf16(float f) : bits(float2bfloat(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;
+}

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

CASE(OverheadType::kU8, OverheadType::kU64, PrimaryType::kF32, uint8_t,
        uint64_t, float);
   CASE(OverheadType::kU8, OverheadType::kU32, PrimaryType::kF32, uint8_t,
        uint32_t, float);
   CASE(OverheadType::kU8, OverheadType::kU16, PrimaryType::kF32, uint8_t,
        uint16_t, float);
   CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kF32, uint8_t,
        uint8_t, float);
 
+  // Two-byte floats with both overheads of the same type.
+  CASE_SECSAME(OverheadType::kU64, PrimaryType::kF16, uint64_t, f16);
+  CASE_SECSAME(OverheadType::kU64, PrimaryType::kBF16, uint64_t, bf16);
+  CASE_SECSAME(OverheadType::kU32, PrimaryType::kF16, uint32_t, f16);
+  CASE_SECSAME(OverheadType::kU32, PrimaryType::kBF16, uint32_t, bf16);
+  CASE_SECSAME(OverheadType::kU16, PrimaryType::kF16, uint16_t, f16);
+  CASE_SECSAME(OverheadType::kU16, PrimaryType::kBF16, uint16_t, bf16);
+  CASE_SECSAME(OverheadType::kU8, PrimaryType::kF16, uint8_t, f16);
+  CASE_SECSAME(OverheadType::kU8, PrimaryType::kBF16, uint8_t, bf16);
+
   // Integral matrices with both overheads of the same type.
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI64, uint64_t, int64_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI32, uint64_t, int32_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI16, uint64_t, int16_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI8, uint64_t, int8_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI64, uint32_t, int64_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI32, uint32_t, int32_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI16, uint32_t, int16_t);

mlir/test/Dialect/SparseTensor/conversion_sparse2dense.mlir

 //   CHECK-DAG: memref.store %[[Attr0]], %[[AttrsS]][%[[I0]]] : memref<1xi8>
 //   CHECK-DAG: %[[SizesS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[SizesD:.*]] = memref.cast %[[SizesS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I13]], %[[SizesS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[PermS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[PermD:.*]] = memref.cast %[[PermS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
-//   CHECK-DAG: %[[ElemTp:.*]] = arith.constant 4 : i32
-//   CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
-//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32
+//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
 //   CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<13xi32>
 //   CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) outs(%[[M]] : memref<13xi32>)
 //       CHECK: scf.while : () -> () {
 //       CHECK:   %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
 //       CHECK:   scf.condition(%[[Cond]])

 //   CHECK-DAG: %[[SizesS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[SizesD:.*]] = memref.cast %[[SizesS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[SizeI0:.*]] = call @sparseDimSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[SizeI0]], %[[SizesS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[PermS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[PermD:.*]] = memref.cast %[[PermS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
-//   CHECK-DAG: %[[ElemTp:.*]] = arith.constant 4 : i32
-//   CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
-//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32
+//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
 //   CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?xi32>
 //   CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) outs(%[[M]] : memref<?xi32>)
 //       CHECK: scf.while : () -> () {
 //       CHECK:   %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
 //       CHECK:   scf.condition(%[[Cond]])

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.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<?xf16, #SparseVector>,
+                        %argb: tensor<?xf16, #SparseVector>) -> tensor<?xf16, #DenseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xf16, #SparseVector>
+    %xv = bufferization.alloc_tensor (%d) : tensor<?xf16, #DenseVector>
+    %0 = linalg.generic #trait_vec_op
+       ins(%arga, %argb: tensor<?xf16, #SparseVector>, tensor<?xf16, #SparseVector>)
+        outs(%xv: tensor<?xf16, #DenseVector>) {
+        ^bb(%a: f16, %b: f16, %x: f16):
+          %1 = sparse_tensor.binary %a, %b : f16, f16 to f16
+            overlap={
+              ^bb0(%a0: f16, %b0: f16):
+                %cmp = arith.cmpf "olt", %a0, %b0 : f16
+                %2 = arith.select %cmp, %a0, %b0: f16
+                sparse_tensor.yield %2 : f16
+            }
+            left=identity
+            right=identity
+          linalg.yield %1 : f16
+    } -> tensor<?xf16, #DenseVector>
+    return %0 : tensor<?xf16, #DenseVector>
+  }
+
+  // Dumps a dense vector of type f16.
+  func.func @dump_vec(%arg0: tensor<?xf16, #DenseVector>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f16
+    %0 = sparse_tensor.values %arg0 : tensor<?xf16, #DenseVector> to memref<?xf16>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf16>, vector<32xf16>
+    %f1 = arith.extf %1: vector<32xf16> 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<32xf16>
+    %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<32xf16>
+    %sv1 = sparse_tensor.convert %v1 : tensor<32xf16> to tensor<?xf16, #SparseVector>
+    %sv2 = sparse_tensor.convert %v2 : tensor<32xf16> to tensor<?xf16, #SparseVector>
+
+    // Call the sparse vector kernel.
+    %0 = call @vector_min(%sv1, %sv2)
+       : (tensor<?xf16, #SparseVector>,
+          tensor<?xf16, #SparseVector>) -> tensor<?xf16, #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<?xf16, #DenseVector>) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sv1 : tensor<?xf16, #SparseVector>
+    sparse_tensor.release %sv2 : tensor<?xf16, #SparseVector>
+    sparse_tensor.release %0 : tensor<?xf16, #DenseVector>
+    return
+  }
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.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?xf16, #SparseMatrix>,
+                          %argx: tensor<f16> {linalg.inplaceable = true}) -> tensor<f16> {
+    %0 = linalg.generic #trait_sum_reduce
+      ins(%arga: tensor<?x?xf16, #SparseMatrix>)
+      outs(%argx: tensor<f16>) {
+      ^bb(%a: f16, %x: f16):
+        %0 = arith.addf %x, %a : f16
+        linalg.yield %0 : f16
+    } -> tensor<f16>
+    return %0 : tensor<f16>
+  }
+
+  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<3x4xf16>
+    %a = sparse_tensor.convert %d : tensor<3x4xf16> to tensor<?x?xf16, #SparseMatrix>
+
+    %d0 = arith.constant 0.0 : f16
+    // Setup memory for a single reduction scalar,
+    // initialized to zero.
+    %xdata = memref.alloc() : memref<f16>
+    memref.store %d0, %xdata[] : memref<f16>
+    %x = bufferization.to_tensor %xdata : memref<f16>
+
+    // Call the kernel.
+    %0 = call @kernel_sum_reduce(%a, %x)
+      : (tensor<?x?xf16, #SparseMatrix>, tensor<f16>) -> tensor<f16>
+
+    // Print the result for verification.
+    //
+    // CHECK: 13.5
+    //
+    %m = bufferization.to_memref %0 : memref<f16>
+    %v = memref.load %m[] : memref<f16>
+    %vf = arith.extf %v: f16 to f32
+    vector.print %vf : f32
+
+    // Release the resources.
+    memref.dealloc %xdata : memref<f16>
+    sparse_tensor.release %a : tensor<?x?xf16, #SparseMatrix>
+
+    return
+  }
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "mlir_c_runner_utils",
     srcs = [
         "lib/ExecutionEngine/CRunnerUtils.cpp",
+        "lib/ExecutionEngine/Float16bits.cpp",
         "lib/ExecutionEngine/SparseTensorUtils.cpp",
     ],
     hdrs = [
         "include/mlir/ExecutionEngine/CRunnerUtils.h",
+        "include/mlir/ExecutionEngine/Float16bits.h",
         "include/mlir/ExecutionEngine/Msan.h",
         "include/mlir/ExecutionEngine/SparseTensorUtils.h",
     ],
     includes = ["include"],
 )
 
 cc_library(
     name = "mlir_async_runtime_api",