Add SPV_KHR_bfloat16 support (KhronosGroup#6057)

* Add SPV_KHR_bfloat16 support

* Update DEPS to include SPIRV-Headers with bfloat16 support

* Fix unit test errors and format

* Add validation to invalid uses of bfloat16

* Add tests

* Roll back to previous commit

* Fix build error

* Add FPEncoding for opt::analysis::Float

* Address the comments

* Fix build error

* format

---------

Co-authored-by: Stu Smith <[email protected]>
Co-authored-by: David Neto <[email protected]>

Author: 3 people

Android.mk

@@ -76,7 +76,8 @@ SPVTOOLS_SRC_FILES := \
 		source/val/validate_scopes.cpp \
 		source/val/validate_small_type_uses.cpp \
 		source/val/validate_tensor_layout.cpp \
-		source/val/validate_type.cpp
+		source/val/validate_type.cpp\
+		source/val/validate_invalid_type.cpp
 
 SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/aggressive_dead_code_elim_pass.cpp \

BUILD.gn

@@ -559,6 +559,7 @@ static_library("spvtools_val") {
     "source/val/validate_small_type_uses.cpp",
     "source/val/validate_tensor_layout.cpp",
     "source/val/validate_type.cpp",
+    "source/val/validate_invalid_type.cpp",
     "source/val/validation_state.cpp",
     "source/val/validation_state.h",
   ]

source/CMakeLists.txt

@@ -336,6 +336,7 @@ set(SPIRV_SOURCES
   ${CMAKE_CURRENT_SOURCE_DIR}/val/validate_small_type_uses.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/val/validate_tensor_layout.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/val/validate_type.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/val/validate_invalid_type.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/val/decoration.h
   ${CMAKE_CURRENT_SOURCE_DIR}/val/basic_block.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/val/construct.cpp

source/name_mapper.cpp

@@ -211,7 +211,12 @@ spv_result_t FriendlyNameMapper::ParseInstruction(
     } break;
     case spv::Op::OpTypeFloat: {
       const auto bit_width = inst.words[2];
-      // TODO: Handle optional fpencoding enum once actually used.
+      if (inst.num_words > 3) {
+        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::BFloat16KHR) {
+          SaveName(result_id, "bfloat16");
+          break;
+        }
+      }
       switch (bit_width) {
         case 16:
           SaveName(result_id, "half");

source/opt/type_manager.cpp

@@ -792,9 +792,13 @@ Type* TypeManager::RecordIfTypeDefinition(const Instruction& inst) {
       type = new Integer(inst.GetSingleWordInOperand(0),
                          inst.GetSingleWordInOperand(1));
       break;
-    case spv::Op::OpTypeFloat:
-      type = new Float(inst.GetSingleWordInOperand(0));
-      break;
+    case spv::Op::OpTypeFloat: {
+      const spv::FPEncoding encoding =
+          inst.NumInOperands() > 1
+              ? static_cast<spv::FPEncoding>(inst.GetSingleWordInOperand(1))
+              : spv::FPEncoding::Max;
+      type = new Float(inst.GetSingleWordInOperand(0), encoding);
+    } break;
     case spv::Op::OpTypeVector:
       type = new Vector(GetType(inst.GetSingleWordInOperand(0)),
                         inst.GetSingleWordInOperand(1));

source/opt/types.cpp

@@ -309,17 +309,26 @@ size_t Integer::ComputeExtraStateHash(size_t hash, SeenTypes*) const {
 
 bool Float::IsSameImpl(const Type* that, IsSameCache*) const {
   const Float* ft = that->AsFloat();
-  return ft && width_ == ft->width_ && HasSameDecorations(that);
+  return ft && width_ == ft->width_ && encoding_ == ft->encoding_ &&
+         HasSameDecorations(that);
 }
 
 std::string Float::str() const {
   std::ostringstream oss;
-  oss << "float" << width_;
+  switch (encoding_) {
+    case spv::FPEncoding::BFloat16KHR:
+      assert(width_ == 16);
+      oss << "bfloat16";
+      break;
+    default:
+      oss << "float" << width_;
+      break;
+  }
   return oss.str();
 }
 
 size_t Float::ComputeExtraStateHash(size_t hash, SeenTypes*) const {
-  return hash_combine(hash, width_);
+  return hash_combine(hash, width_, encoding_);
 }
 
 Vector::Vector(const Type* type, uint32_t count)

source/opt/types.h

@@ -182,9 +182,9 @@ class Type {
   // non-composite type.
   uint64_t NumberOfComponents() const;
 
-// A bunch of methods for casting this type to a given type. Returns this if the
-// cast can be done, nullptr otherwise.
-// clang-format off
+  // A bunch of methods for casting this type to a given type. Returns this if
+  // the cast can be done, nullptr otherwise.
+  // clang-format off
 #define DeclareCastMethod(target)                  \
   virtual target* As##target() { return nullptr; } \
   virtual const target* As##target() const { return nullptr; }
@@ -267,21 +267,24 @@ class Integer : public Type {
 
 class Float : public Type {
  public:
-  Float(uint32_t w) : Type(kFloat), width_(w) {}
+  Float(uint32_t w, spv::FPEncoding encoding = spv::FPEncoding::Max)
+      : Type(kFloat), width_(w), encoding_(encoding) {}
   Float(const Float&) = default;
 
   std::string str() const override;
 
   Float* AsFloat() override { return this; }
   const Float* AsFloat() const override { return this; }
   uint32_t width() const { return width_; }
+  spv::FPEncoding encoding() const { return encoding_; }
 
   size_t ComputeExtraStateHash(size_t hash, SeenTypes* seen) const override;
 
  private:
   bool IsSameImpl(const Type* that, IsSameCache*) const override;
 
-  uint32_t width_;  // bit width
+  uint32_t width_;            // bit width
+  spv::FPEncoding encoding_;  // FPEncoding
 };
 
 class Vector : public Type {

source/val/validate.cpp

@@ -367,6 +367,7 @@ spv_result_t ValidateBinaryUsingContextAndValidationState(
     if (auto error = RayReorderNVPass(*vstate, &instruction)) return error;
     if (auto error = MeshShadingPass(*vstate, &instruction)) return error;
     if (auto error = TensorLayoutPass(*vstate, &instruction)) return error;
+    if (auto error = InvalidTypePass(*vstate, &instruction)) return error;
   }
 
   // Validate the preconditions involving adjacent instructions. e.g.

source/val/validate.h

@@ -223,6 +223,9 @@ spv_result_t RayReorderNVPass(ValidationState_t& _, const Instruction* inst);
 /// Validates correctness of mesh shading instructions.
 spv_result_t MeshShadingPass(ValidationState_t& _, const Instruction* inst);
 
+/// Validates correctness of certain special type instructions.
+spv_result_t InvalidTypePass(ValidationState_t& _, const Instruction* inst);
+
 /// Calculates the reachability of basic blocks.
 void ReachabilityPass(ValidationState_t& _);
 

source/val/validate_arithmetics.cpp

@@ -224,6 +224,14 @@ spv_result_t ArithmeticsPass(ValidationState_t& _, const Instruction* inst) {
                << "Expected float scalar type as Result Type: "
                << spvOpcodeString(opcode);
 
+      if (_.IsBfloat16ScalarType(result_type)) {
+        if (!_.HasCapability(spv::Capability::BFloat16DotProductKHR)) {
+          return _.diag(SPV_ERROR_INVALID_DATA, inst)
+                 << "OpDot Result Type <id> " << _.getIdName(result_type)
+                 << "requires BFloat16DotProductKHR be declared.";
+        }
+      }
+
       uint32_t first_vector_num_components = 0;
 
       for (size_t operand_index = 2; operand_index < inst->operands().size();