[mlir] Add support for bf16 to StandardToLLVM conversion

[dcaballe]

Summary

This patch adds support to convert std bf16 types to LLVM bfloat types.

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

public:
   using Base::Base;
 
   static bool kindof(unsigned kind) { return kind == LLVM_TYPE; }
 
   LLVMDialect &getDialect();
   llvm::Type *getUnderlyingType() const;
 
   /// Utilities to identify types.
+  bool isBFloatTy() { return getUnderlyingType()->isBFloatTy(); }
   bool isHalfTy() { return getUnderlyingType()->isHalfTy(); }
   bool isFloatTy() { return getUnderlyingType()->isFloatTy(); }
   bool isDoubleTy() { return getUnderlyingType()->isDoubleTy(); }
   bool isIntegerTy() { return getUnderlyingType()->isIntegerTy(); }
   bool isIntegerTy(unsigned bitwidth) {
     return getUnderlyingType()->isIntegerTy(bitwidth);
   }
 
public:
   /// Struct type utilities.
   LLVMType getStructElementType(unsigned i);
   unsigned getStructNumElements();
   bool isStructTy();
 
   /// Utilities used to generate floating point types.
   static LLVMType getDoubleTy(LLVMDialect *dialect);
   static LLVMType getFloatTy(LLVMDialect *dialect);
+  static LLVMType getBFloatTy(LLVMDialect *dialect);
   static LLVMType getHalfTy(LLVMDialect *dialect);
   static LLVMType getFP128Ty(LLVMDialect *dialect);
   static LLVMType getX86_FP80Ty(LLVMDialect *dialect);
 
   /// Utilities used to generate integer types.
   static LLVMType getIntNTy(LLVMDialect *dialect, unsigned numBits);
   static LLVMType getInt1Ty(LLVMDialect *dialect) {
     return getIntNTy(dialect, /*numBits=*/1);

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

Type LLVMTypeConverter::convertFloatType(FloatType type) {
   switch (type.getKind()) {
   case mlir::StandardTypes::F32:
     return LLVM::LLVMType::getFloatTy(llvmDialect);
   case mlir::StandardTypes::F64:
     return LLVM::LLVMType::getDoubleTy(llvmDialect);
   case mlir::StandardTypes::F16:
     return LLVM::LLVMType::getHalfTy(llvmDialect);
   case mlir::StandardTypes::BF16: {
-    auto *mlirContext = llvmDialect->getContext();
-    return emitError(UnknownLoc::get(mlirContext), "unsupported type: BF16"),
-           Type();
+    return LLVM::LLVMType::getBFloatTy(llvmDialect);
   }
   default:
     llvm_unreachable("non-float type in convertFloatType");
   }
 }
 
 // Convert a `ComplexType` to an LLVM type. The result is a complex number
 // struct with entries for the

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 // Verifier for LLVM::DialectCastOp.
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(DialectCastOp op) {
   auto verifyMLIRCastType = [&op](Type type) -> LogicalResult {
     if (auto llvmType = type.dyn_cast<LLVM::LLVMType>()) {
       if (llvmType.isVectorTy())
         llvmType = llvmType.getVectorElementType();
-      if (llvmType.isIntegerTy() || llvmType.isHalfTy() ||
-          llvmType.isFloatTy() || llvmType.isDoubleTy()) {
+      if (llvmType.isIntegerTy() || llvmType.isBFloatTy() ||
+          llvmType.isHalfTy() || llvmType.isFloatTy() ||
+          llvmType.isDoubleTy()) {
         return success();
       }
       return op.emitOpError("type must be non-index integer types, float "
                             "types, or vector of mentioned types.");
     }
     if (auto vectorType = type.dyn_cast<VectorType>()) {
       if (vectorType.getShape().size() > 1)
         return op.emitOpError("only 1-d vector is allowed");
if (resType != valType)
     return op.emitOpError(
         "expected LLVM IR result type to match type for operand #1");
   if (op.bin_op() == AtomicBinOp::fadd || op.bin_op() == AtomicBinOp::fsub) {
     if (!valType.getUnderlyingType()->isFloatingPointTy())
       return op.emitOpError("expected LLVM IR floating point type");
   } else if (op.bin_op() == AtomicBinOp::xchg) {
     if (!valType.isIntegerTy(8) && !valType.isIntegerTy(16) &&
         !valType.isIntegerTy(32) && !valType.isIntegerTy(64) &&
-        !valType.isHalfTy() && !valType.isFloatTy() && !valType.isDoubleTy())
+        !valType.isBFloatTy() && !valType.isHalfTy() && !valType.isFloatTy() &&
+        !valType.isDoubleTy())
       return op.emitOpError("unexpected LLVM IR type for 'xchg' bin_op");
   } else {
     if (!valType.isIntegerTy(8) && !valType.isIntegerTy(16) &&
         !valType.isIntegerTy(32) && !valType.isIntegerTy(64))
       return op.emitOpError("expected LLVM IR integer type");
   }
   return success();
 }
if (!ptrType.isPointerTy())
     return op.emitOpError("expected LLVM IR pointer type for operand #0");
   auto cmpType = op.cmp().getType().cast<LLVM::LLVMType>();
   auto valType = op.val().getType().cast<LLVM::LLVMType>();
   if (cmpType != ptrType.getPointerElementTy() || cmpType != valType)
     return op.emitOpError("expected LLVM IR element type for operand #0 to "
                           "match type for all other operands");
   if (!valType.isPointerTy() && !valType.isIntegerTy(8) &&
       !valType.isIntegerTy(16) && !valType.isIntegerTy(32) &&
-      !valType.isIntegerTy(64) && !valType.isHalfTy() && !valType.isFloatTy() &&
-      !valType.isDoubleTy())
+      !valType.isIntegerTy(64) && !valType.isBFloatTy() &&
+      !valType.isHalfTy() && !valType.isFloatTy() && !valType.isDoubleTy())
     return op.emitOpError("unexpected LLVM IR type");
   if (op.success_ordering() < AtomicOrdering::monotonic ||
       op.failure_ordering() < AtomicOrdering::monotonic)
     return op.emitOpError("ordering must be at least 'monotonic'");
   if (op.failure_ordering() == AtomicOrdering::release ||
       op.failure_ordering() == AtomicOrdering::acq_rel)
     return op.emitOpError("failure ordering cannot be 'release' or 'acq_rel'");
   return success();
struct LLVMDialectImpl {
   LLVMDialectImpl() : module("LLVMDialectModule", llvmContext) {}
 
   llvm::LLVMContext llvmContext;
   llvm::Module module;
 
   /// A set of LLVMTypes that are cached on construction to avoid any lookups or
   /// locking.
   LLVMType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
-  LLVMType doubleTy, floatTy, halfTy, fp128Ty, x86_fp80Ty;
+  LLVMType doubleTy, floatTy, bfloatTy, halfTy, fp128Ty, x86_fp80Ty;
   LLVMType voidTy;
 
   /// A smart mutex to lock access to the llvm context. Unlike MLIR, LLVM is not
   /// multi-threaded and requires locked access to prevent race conditions.
   llvm::sys::SmartMutex<true> mutex;
 };
 } // end namespace detail
 } // end namespace LLVM
#include "mlir/Dialect/LLVMIR/LLVMOps.cpp.inc"
   impl->int8Ty = LLVMType::get(context, llvm::Type::getInt8Ty(llvmContext));
   impl->int16Ty = LLVMType::get(context, llvm::Type::getInt16Ty(llvmContext));
   impl->int32Ty = LLVMType::get(context, llvm::Type::getInt32Ty(llvmContext));
   impl->int64Ty = LLVMType::get(context, llvm::Type::getInt64Ty(llvmContext));
   impl->int128Ty = LLVMType::get(context, llvm::Type::getInt128Ty(llvmContext));
   /// Float Types.
   impl->doubleTy = LLVMType::get(context, llvm::Type::getDoubleTy(llvmContext));
   impl->floatTy = LLVMType::get(context, llvm::Type::getFloatTy(llvmContext));
+  impl->bfloatTy = LLVMType::get(context, llvm::Type::getBFloatTy(llvmContext));
   impl->halfTy = LLVMType::get(context, llvm::Type::getHalfTy(llvmContext));
   impl->fp128Ty = LLVMType::get(context, llvm::Type::getFP128Ty(llvmContext));
   impl->x86_fp80Ty =
       LLVMType::get(context, llvm::Type::getX86_FP80Ty(llvmContext));
   /// Other Types.
   impl->voidTy = LLVMType::get(context, llvm::Type::getVoidTy(llvmContext));
 }
 

 
 /// Utilities used to generate floating point types.
 LLVMType LLVMType::getDoubleTy(LLVMDialect *dialect) {
   return dialect->impl->doubleTy;
 }
 LLVMType LLVMType::getFloatTy(LLVMDialect *dialect) {
   return dialect->impl->floatTy;
 }
+LLVMType LLVMType::getBFloatTy(LLVMDialect *dialect) {
+  return dialect->impl->bfloatTy;
+}
 LLVMType LLVMType::getHalfTy(LLVMDialect *dialect) {
   return dialect->impl->halfTy;
 }
 LLVMType LLVMType::getFP128Ty(LLVMDialect *dialect) {
   return dialect->impl->fp128Ty;
 }
 LLVMType LLVMType::getX86_FP80Ty(LLVMDialect *dialect) {
   return dialect->impl->x86_fp80Ty;

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 
 // CHECK-LABEL: func @mlir_cast_from_llvm
 // CHECK-SAME: %[[ARG:.*]]:
 func @mlir_cast_from_llvm(%0 : !llvm.half) -> f16 {
   %1 = llvm.mlir.cast %0 : !llvm.half to f16
   // CHECK-NEXT: llvm.return %[[ARG]]
   return %1 : f16
 }
+
+// -----
+
+// CHECK-LABEL: func @bfloat
+// CHECK-SAME: !llvm.bfloat) -> !llvm.bfloat
+func @bfloat(%arg0: bf16) -> bf16 {
+  return %arg0 : bf16
+}
+// CHECK-NEXT: return %{{.*}} : !llvm.bfloat