maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][sparse] Add F16 and BF16. 

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.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D127010
This commit is contained in:
bixia1 2022-06-07 16:07:13 -07:00
parent 9a76337fee
commit ea8ed5cbcf
10 changed files with 383 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
mlir/include/mlir/ExecutionEngine/Float16bits.h Normal file
View File

@ -0,0 +1,39 @@
//===--- 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_

17
mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h
View File

@ -15,6 +15,7 @@
#define MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
#include "mlir/ExecutionEngine/CRunnerUtils.h"
#include "mlir/ExecutionEngine/Float16bits.h"
#include <cinttypes>
#include <complex>
@ -77,12 +78,14 @@ using complex32 = std::complex<float>;
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
@ -97,6 +100,8 @@ enum class PrimaryType : uint32_t {
#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) \

4
mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
View File

@ -104,6 +104,10 @@ PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
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))

2
mlir/lib/ExecutionEngine/CMakeLists.txt
View File

@ -7,6 +7,7 @@ set(LLVM_OPTIONAL_SOURCES
CudaRuntimeWrappers.cpp
SparseTensorUtils.cpp
ExecutionEngine.cpp
Float16bits.cpp
RocmRuntimeWrappers.cpp
RunnerUtils.cpp
OptUtils.cpp
@ -121,6 +122,7 @@ add_mlir_library(MLIRJitRunner
add_mlir_library(mlir_c_runner_utils
SHARED
CRunnerUtils.cpp
Float16bits.cpp
SparseTensorUtils.cpp
EXCLUDE_FROM_LIBMLIR

143
mlir/lib/ExecutionEngine/Float16bits.cpp Normal file
View File

@ -0,0 +1,143 @@
//===--- 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;
}

10
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
View File

@ -1567,6 +1567,16 @@ _mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
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);

10
mlir/test/Dialect/SparseTensor/conversion_sparse2dense.mlir
View File

@ -28,9 +28,8 @@
// 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>
@ -67,9 +66,8 @@ func.func @sparse_convert_1d(%arg0: tensor<13xi32, #SparseVector>) -> tensor<13x
// 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>

90
mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir Normal file
View File

@ -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<?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
}
}

78
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir Normal file
View File

@ -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?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
}
}

2
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
View File

@ -6415,10 +6415,12 @@ 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",
],