[HLSL] Allow register annotations to specify only space

Specifying only `space` in a `register` annotation means the compiler should
implicitly assign a register slot to the resource from the provided virtual
register space.

Author: hekota

clang/include/clang/Basic/Attr.td

@@ -4751,20 +4751,25 @@ def HLSLResourceBinding: InheritableAttr {
 
   private:
       RegisterType RegType;
-      unsigned SlotNumber;
+      int SlotNumber; // -1 if the register slot was not specified
       unsigned SpaceNumber;
 
   public:
-      void setBinding(RegisterType RT, unsigned SlotNum, unsigned SpaceNum) {
+      void setBinding(RegisterType RT, int SlotNum, unsigned SpaceNum) {
         RegType = RT;
         SlotNumber = SlotNum;
         SpaceNumber = SpaceNum;
       }
+      bool isImplicit() const {
+        return SlotNumber < 0;
+      }
       RegisterType getRegisterType() const {
+        assert(!isImplicit() && "binding does not have register slot");
         return RegType;
       }
       unsigned getSlotNumber() const {
-        return SlotNumber;
+        assert(!isImplicit() && "binding does not have register slot");
+        return (unsigned)SlotNumber;
       }
       unsigned getSpaceNumber() const {
         return SpaceNumber;

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -261,7 +261,7 @@ void CGHLSLRuntime::addBuffer(const HLSLBufferDecl *BufDecl) {
       BufDecl->getAttr<HLSLResourceBindingAttr>();
   // FIXME: handle implicit binding if no binding attribute is found
   // (llvm/llvm-project#110722)
-  if (RBA)
+  if (RBA && !RBA->isImplicit())
     initializeBufferFromBinding(CGM, BufGV, RBA->getSlotNumber(),
                                 RBA->getSpaceNumber());
 }

clang/lib/Sema/SemaHLSL.cpp

@@ -1529,72 +1529,82 @@ void SemaHLSL::handleResourceBindingAttr(Decl *TheDecl, const ParsedAttr &AL) {
                                     diag::err_incomplete_type))
       return;
   }
-  StringRef Space = "space0";
+
   StringRef Slot = "";
+  StringRef Space = "";
+  SourceLocation SlotLoc, SpaceLoc;
 
   if (!AL.isArgIdent(0)) {
     Diag(AL.getLoc(), diag::err_attribute_argument_type)
         << AL << AANT_ArgumentIdentifier;
     return;
   }
-
   IdentifierLoc *Loc = AL.getArgAsIdent(0);
-  StringRef Str = Loc->Ident->getName();
-  SourceLocation ArgLoc = Loc->Loc;
 
-  SourceLocation SpaceArgLoc;
-  bool SpecifiedSpace = false;
   if (AL.getNumArgs() == 2) {
-    SpecifiedSpace = true;
-    Slot = Str;
+    Slot = Loc->Ident->getName();
+    SlotLoc = Loc->Loc;
+
     if (!AL.isArgIdent(1)) {
       Diag(AL.getLoc(), diag::err_attribute_argument_type)
           << AL << AANT_ArgumentIdentifier;
       return;
     }
 
-    IdentifierLoc *Loc = AL.getArgAsIdent(1);
+    Loc = AL.getArgAsIdent(1);
     Space = Loc->Ident->getName();
-    SpaceArgLoc = Loc->Loc;
+    SpaceLoc = Loc->Loc;
   } else {
-    Slot = Str;
+    StringRef Str = Loc->Ident->getName();
+    if (Str.starts_with("space")) {
+      Space = Str;
+      SpaceLoc = Loc->Loc;
+    } else {
+      Slot = Str;
+      SlotLoc = Loc->Loc;
+      Space = "space0";
+    }
   }
 
-  RegisterType RegType;
-  unsigned SlotNum = 0;
+  RegisterType RegType = RegisterType::SRV;
+  int SlotNum = -1;
   unsigned SpaceNum = 0;
 
-  // Validate.
+  // Validate slot
   if (!Slot.empty()) {
     if (!convertToRegisterType(Slot, &RegType)) {
-      Diag(ArgLoc, diag::err_hlsl_binding_type_invalid) << Slot.substr(0, 1);
+      Diag(SlotLoc, diag::err_hlsl_binding_type_invalid) << Slot.substr(0, 1);
       return;
     }
     if (RegType == RegisterType::I) {
-      Diag(ArgLoc, diag::warn_hlsl_deprecated_register_type_i);
+      Diag(SlotLoc, diag::warn_hlsl_deprecated_register_type_i);
       return;
     }
-
     StringRef SlotNumStr = Slot.substr(1);
     if (SlotNumStr.getAsInteger(10, SlotNum)) {
-      Diag(ArgLoc, diag::err_hlsl_unsupported_register_number);
+      Diag(SlotLoc, diag::err_hlsl_unsupported_register_number);
       return;
     }
   }
 
-  if (!Space.starts_with("space")) {
-    Diag(SpaceArgLoc, diag::err_hlsl_expected_space) << Space;
-    return;
-  }
-  StringRef SpaceNumStr = Space.substr(5);
-  if (SpaceNumStr.getAsInteger(10, SpaceNum)) {
-    Diag(SpaceArgLoc, diag::err_hlsl_expected_space) << Space;
-    return;
+  // Validate space
+  if (!Space.empty()) {
+    if (!Space.starts_with("space")) {
+      Diag(SpaceLoc, diag::err_hlsl_expected_space) << Space;
+      return;
+    }
+    StringRef SpaceNumStr = Space.substr(5);
+    if (SpaceNumStr.getAsInteger(10, SpaceNum)) {
+      Diag(SpaceLoc, diag::err_hlsl_expected_space) << Space;
+      return;
+    }
   }
 
-  if (!DiagnoseHLSLRegisterAttribute(SemaRef, ArgLoc, TheDecl, RegType,
-                                     SpecifiedSpace))
-    return;
+  // If we have slot, diagnose it is the right register type for the decl
+  if (SlotNum >= 0)
+    if (!DiagnoseHLSLRegisterAttribute(SemaRef, SlotLoc, TheDecl, RegType,
+                                       !SpaceLoc.isInvalid()))
+      return;
 
   HLSLResourceBindingAttr *NewAttr =
       HLSLResourceBindingAttr::Create(getASTContext(), Slot, Space, AL);
@@ -1967,7 +1977,7 @@ void SemaHLSL::ActOnEndOfTranslationUnit(TranslationUnitDecl *TU) {
     for (const Decl *VD : DefaultCBufferDecls) {
       const HLSLResourceBindingAttr *RBA =
           VD->getAttr<HLSLResourceBindingAttr>();
-      if (RBA &&
+      if (RBA && !RBA->isImplicit() &&
           RBA->getRegisterType() == HLSLResourceBindingAttr::RegisterType::C) {
         DefaultCBuffer->setHasValidPackoffset(true);
         break;
@@ -3227,7 +3237,7 @@ static bool initVarDeclWithCtor(Sema &S, VarDecl *VD,
 
 static bool initGlobalResourceDecl(Sema &S, VarDecl *VD) {
   HLSLResourceBindingAttr *RBA = VD->getAttr<HLSLResourceBindingAttr>();
-  if (!RBA)
+  if (!RBA || RBA->isImplicit())
     // FIXME: add support for implicit binding (llvm/llvm-project#110722)
     return false;
 
@@ -3310,7 +3320,7 @@ void SemaHLSL::processExplicitBindingsOnDecl(VarDecl *VD) {
 
   for (Attr *A : VD->attrs()) {
     HLSLResourceBindingAttr *RBA = dyn_cast<HLSLResourceBindingAttr>(A);
-    if (!RBA)
+    if (!RBA || RBA->isImplicit())
       continue;
 
     RegisterType RT = RBA->getRegisterType();

clang/test/AST/HLSL/resource_binding_attr.hlsl

@@ -30,6 +30,10 @@ RWBuffer<float> UAV : register(u3);
 // CHECK: HLSLResourceBindingAttr {{.*}} "u4" "space0"
 RWBuffer<float> UAV1 : register(u2), UAV2 : register(u4);
 
+// CHECK: VarDecl {{.*}} UAV3 'RWBuffer<float>':'hlsl::RWBuffer<float>'
+// CHECK: HLSLResourceBindingAttr {{.*}} "" "space5"
+RWBuffer<float> UAV3 : register(space5);
+
 //
 // Default constants ($Globals) layout annotations
 

clang/test/SemaHLSL/resource_binding_attr_error.hlsl

@@ -22,8 +22,8 @@ cbuffer c : register(bf, s2) {
 // expected-error@+1 {{expected identifier}}
 cbuffer A : register() {}
 
-// expected-error@+1 {{register number should be an integer}}
-cbuffer B : register(space1) {}
+// expected-error@+1 {{invalid space specifier 'space' used; expected 'space' followed by an integer, like space1}}
+cbuffer B : register(space) {}
 
 // expected-error@+1 {{wrong argument format for hlsl attribute, use b2 instead}}
 cbuffer C : register(b 2) {}

test.ll

@@ -0,0 +1,7 @@
+ByteAddressBuffer Buffer0: register(t0);
+RWBuffer<float> Buf : register(u0);
+
+[numthreads(4,1,1)]
+void main() {
+  Buf[0] = 10;
+}