[spirv] Use extension for fp16 in GLSL extended instruction set. (#2151)

tools/clang/include/clang/SPIRV/SpirvType.h

@@ -68,7 +68,8 @@ public:
   static bool isSubpassInput(const SpirvType *);
   static bool isSubpassInputMS(const SpirvType *);
   static bool isResourceType(const SpirvType *);
-  static bool isOrContains16BitType(const SpirvType *);
+  template <class T, unsigned int Bitwidth = 0>
+  static bool isOrContainsType(const SpirvType *);
 
 protected:
   SpirvType(Kind k, llvm::StringRef name = "") : kind(k), debugName(name) {}
@@ -525,6 +526,44 @@ private:
   llvm::SmallVector<QualType, 8> paramTypes;
 };
 
+//
+// Function Definition for templated functions
+//
+
+template <class T, unsigned int Bitwidth>
+bool SpirvType::isOrContainsType(const SpirvType *type) {
+  if (isa<T>(type)) {
+    if (Bitwidth == 0)
+      // No specific bitwidth was asked for.
+      return true;
+    else
+      // We want to make sure it is a numberical type of a specific bitwidth.
+      return isa<NumericalType>(type) &&
+             llvm::cast<NumericalType>(type)->getBitwidth() == Bitwidth;
+  }
+
+  if (const auto *vecType = dyn_cast<VectorType>(type))
+    return isOrContainsType<T, Bitwidth>(vecType->getElementType());
+  if (const auto *matType = dyn_cast<MatrixType>(type))
+    return isOrContainsType<T, Bitwidth>(matType->getElementType());
+  if (const auto *arrType = dyn_cast<ArrayType>(type))
+    return isOrContainsType<T, Bitwidth>(arrType->getElementType());
+  if (const auto *pointerType = dyn_cast<SpirvPointerType>(type))
+    return isOrContainsType<T, Bitwidth>(pointerType->getPointeeType());
+  if (const auto *raType = dyn_cast<RuntimeArrayType>(type))
+    return isOrContainsType<T, Bitwidth>(raType->getElementType());
+  if (const auto *imgType = dyn_cast<ImageType>(type))
+    return isOrContainsType<T, Bitwidth>(imgType->getSampledType());
+  if (const auto *sampledImageType = dyn_cast<SampledImageType>(type))
+    return isOrContainsType<T, Bitwidth>(sampledImageType->getImageType());
+  if (const auto *structType = dyn_cast<StructType>(type))
+    for (auto &field : structType->getFields())
+      if (isOrContainsType<T, Bitwidth>(field.type))
+        return true;
+
+  return false;
+}
+
 } // end namespace spirv
 } // end namespace clang
 

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -64,7 +64,6 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
     case 16: {
       // Usage of a 16-bit float type.
       addCapability(spv::Capability::Float16);
-      addExtension(Extension::AMD_gpu_shader_half_float, "16-bit float", loc);
 
       // Usage of a 16-bit float type as stage I/O.
       if (sc == spv::StorageClass::Input || sc == spv::StorageClass::Output) {
@@ -180,7 +179,7 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
   }
   // Struct type
   else if (const auto *structType = dyn_cast<StructType>(type)) {
-    if (SpirvType::isOrContains16BitType(structType)) {
+    if (SpirvType::isOrContainsType<NumericalType, 16>(structType)) {
       addExtension(Extension::KHR_16bit_storage, "16-bit types in resource",
                    loc);
       if (sc == spv::StorageClass::PushConstant) {
@@ -504,5 +503,47 @@ bool CapabilityVisitor::visit(SpirvExecutionMode *execMode) {
   return true;
 }
 
+bool CapabilityVisitor::visit(SpirvExtInst *instr) {
+  // OpExtInst using the GLSL extended instruction allows only 32-bit types by
+  // default. The AMD_gpu_shader_half_float extension adds support for 16-bit
+  // floating-point component types for the following instructions described in
+  // the GLSL.std.450 extended instruction set:
+  // Acos, Acosh, Asin, Asinh, Atan2, Atanh, Atan, Cos, Cosh, Degrees, Exp,
+  // Exp2, InterpolateAtCentroid, InterpolateAtSample, InterpolateAtOffset, Log,
+  // Log2, Pow, Radians, Sin, Sinh, Tan, Tanh
+  if (SpirvType::isOrContainsType<FloatType, 16>(instr->getResultType()))
+    switch (instr->getInstruction()) {
+    case GLSLstd450::GLSLstd450Acos:
+    case GLSLstd450::GLSLstd450Acosh:
+    case GLSLstd450::GLSLstd450Asin:
+    case GLSLstd450::GLSLstd450Asinh:
+    case GLSLstd450::GLSLstd450Atan2:
+    case GLSLstd450::GLSLstd450Atanh:
+    case GLSLstd450::GLSLstd450Atan:
+    case GLSLstd450::GLSLstd450Cos:
+    case GLSLstd450::GLSLstd450Cosh:
+    case GLSLstd450::GLSLstd450Degrees:
+    case GLSLstd450::GLSLstd450Exp:
+    case GLSLstd450::GLSLstd450Exp2:
+    case GLSLstd450::GLSLstd450InterpolateAtCentroid:
+    case GLSLstd450::GLSLstd450InterpolateAtSample:
+    case GLSLstd450::GLSLstd450InterpolateAtOffset:
+    case GLSLstd450::GLSLstd450Log:
+    case GLSLstd450::GLSLstd450Log2:
+    case GLSLstd450::GLSLstd450Pow:
+    case GLSLstd450::GLSLstd450Radians:
+    case GLSLstd450::GLSLstd450Sin:
+    case GLSLstd450::GLSLstd450Sinh:
+    case GLSLstd450::GLSLstd450Tan:
+    case GLSLstd450::GLSLstd450Tanh:
+      addExtension(Extension::AMD_gpu_shader_half_float, "16-bit float",
+                   instr->getSourceLocation());
+    default:
+      break;
+    }
+
+  return visitInstruction(instr);
+}
+
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/CapabilityVisitor.h

@@ -32,6 +32,7 @@ public:
   bool visit(SpirvImageQuery *);
   bool visit(SpirvImageOp *);
   bool visit(SpirvImageSparseTexelsResident *);
+  bool visit(SpirvExtInst *);
 
   /// The "sink" visit function for all instructions.
   ///

tools/clang/lib/SPIRV/SpirvType.cpp

@@ -102,33 +102,6 @@ bool SpirvType::isResourceType(const SpirvType *type) {
   return false;
 }
 
-bool SpirvType::isOrContains16BitType(const SpirvType *type) {
-  if (const auto *numericType = dyn_cast<NumericalType>(type))
-    if (numericType->getBitwidth() == 16)
-      return true;
-
-  if (const auto *vecType = dyn_cast<VectorType>(type))
-    return isOrContains16BitType(vecType->getElementType());
-  if (const auto *matType = dyn_cast<MatrixType>(type))
-    return isOrContains16BitType(matType->getElementType());
-  if (const auto *arrType = dyn_cast<ArrayType>(type))
-    return isOrContains16BitType(arrType->getElementType());
-  if (const auto *pointerType = dyn_cast<SpirvPointerType>(type))
-    return isOrContains16BitType(pointerType->getPointeeType());
-  if (const auto *raType = dyn_cast<RuntimeArrayType>(type))
-    return isOrContains16BitType(raType->getElementType());
-  if (const auto *imgType = dyn_cast<ImageType>(type))
-    return isOrContains16BitType(imgType->getSampledType());
-  if (const auto *sampledImageType = dyn_cast<SampledImageType>(type))
-    return isOrContains16BitType(sampledImageType->getImageType());
-  if (const auto *structType = dyn_cast<StructType>(type))
-    for (auto &field : structType->getFields())
-      if (isOrContains16BitType(field.type))
-        return true;
-
-  return false;
-}
-
 MatrixType::MatrixType(const VectorType *vecType, uint32_t vecCount)
     : SpirvType(TK_Matrix), vectorType(vecType), vectorCount(vecCount) {}
 

tools/clang/test/CodeGenSPIRV/constant.scalar.16bit.enabled.hlsl

@@ -17,7 +17,6 @@
 
 // CHECK: OpCapability Float16
 // CHECK: OpCapability Int16
-// CHECK: OpExtension "SPV_AMD_gpu_shader_half_float"
 
 // CHECK-NOT: OpDecorate %c_half RelaxedPrecision
 // CHECK-NOT: OpDecorate %c_min10float RelaxedPrecision

tools/clang/test/CodeGenSPIRV/extension.GLSLstd450-fp16.allowed.hlsl

@@ -0,0 +1,11 @@
+// Run: %dxc -T ps_6_2 -E main -enable-16bit-types
+
+// The GLSL extended instruction set may only use 16-bit floats if AMD_gpu_shader_half_float extension is used.
+
+// CHECK: OpCapability Float16
+// CHECK: OpExtension "SPV_AMD_gpu_shader_half_float"
+// CHECK: OpExtInstImport "GLSL.std.450"
+void main() {
+  float16_t4 a;
+  float16_t4 result = atan(a);
+}

tools/clang/test/CodeGenSPIRV/extension.GLSLstd450-fp16.not-allowed.hlsl

@@ -0,0 +1,11 @@
+// Run: %dxc -T ps_6_2 -E main -enable-16bit-types -fspv-extension=KHR
+
+// The command line option instructs the compiler to only use core Vulkan features and KHR extensions,
+// and prevents the compiler from using vendor-specific extensions.
+
+// CHECK: 10:23: error: SPIR-V extension 'SPV_AMD_gpu_shader_half_float' required for 16-bit float but not permitted to use
+
+void main() {
+  float16_t4 a;
+  float16_t4 result = atan(a);
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1862,6 +1862,19 @@ TEST_F(FileTest, CapabilityUnique) { runFileTest("capability.unique.hlsl"); }
 // For extension uniqueness
 TEST_F(FileTest, ExtensionUnique) { runFileTest("extension.unique.hlsl"); }
 
+// For vendor-specific extensions
+TEST_F(FileTest, VendorSpecificExtensionAllowed) {
+  // The SPV_AMD_gpu_shader_half_float extension adds support for 16-bit
+  // floating-point component types for a number of instructions in the
+  // GLSL.std.450 extended instruction set.
+  runFileTest("extension.GLSLstd450-fp16.allowed.hlsl");
+}
+TEST_F(FileTest, VendorSpecificExtensionNotAllowed) {
+  // Command line options can entirely prevent the compiler from using
+  // vendor-specific extensions.
+  runFileTest("extension.GLSLstd450-fp16.not-allowed.hlsl", Expect::Failure);
+}
+
 // For RelaxedPrecision decorations
 TEST_F(FileTest, DecorationRelaxedPrecisionBasic) {
   runFileTest("decoration.relaxed-precision.basic.hlsl");