[HLSL] GetDimensions methods for buffer resources

clang/include/clang/Basic/Builtins.td

@@ -4951,6 +4951,18 @@ def HLSLResourceNonUniformIndex : LangBuiltin<"HLSL_LANG"> {
   let Prototype = "uint32_t(uint32_t)";
 }
 
+def HLSLResourceGetDimensions : LangBuiltin<"HLSL_LANG"> {
+  let Spellings = ["__builtin_hlsl_buffer_getdimensions"];
+  let Attributes = [NoThrow];
+  let Prototype = "void(...)";
+}
+
+def HLSLResourceGetStride : LangBuiltin<"HLSL_LANG"> {
+  let Spellings = ["__builtin_hlsl_buffer_getstride"];
+  let Attributes = [NoThrow];
+  let Prototype = "void(...)";
+}
+
 def HLSLAll : LangBuiltin<"HLSL_LANG"> {
   let Spellings = ["__builtin_hlsl_all"];
   let Attributes = [NoThrow, Const];

clang/lib/CodeGen/CGHLSLBuiltins.cpp

@@ -160,6 +160,58 @@ static Value *handleHlslSplitdouble(const CallExpr *E, CodeGenFunction *CGF) {
   return LastInst;
 }
 
+static Value *emitDXILGetDimensions(CodeGenFunction *CGF, Value *Handle,
+                                    Value *MipLevel, LValue *OutArg0,
+                                    LValue *OutArg1 = nullptr,
+                                    LValue *OutArg2 = nullptr,
+                                    LValue *OutArg3 = nullptr) {
+  assert(OutArg0 && "first output argument is required");
+
+  llvm::Type *I32 = CGF->Int32Ty;
+  StructType *RetTy = llvm::StructType::get(I32, I32, I32, I32);
+
+  CallInst *CI = CGF->Builder.CreateIntrinsic(
+      RetTy, llvm::Intrinsic::dx_resource_getdimensions,
+      ArrayRef<Value *>{Handle, MipLevel});
+
+  Value *LastInst = nullptr;
+  unsigned OutArgIndex = 0;
+  for (LValue *OutArg : {OutArg0, OutArg1, OutArg2, OutArg3}) {
+    if (OutArg) {
+      Value *OutArgVal = CGF->Builder.CreateExtractValue(CI, OutArgIndex);
+      LastInst = CGF->Builder.CreateStore(OutArgVal, OutArg->getAddress());
+    }
+    ++OutArgIndex;
+  }
+  assert(LastInst && "no output argument stored?");
+  return LastInst;
+}
+
+static Value *emitBufferGetDimensions(CodeGenFunction *CGF, Value *Handle,
+                                      LValue &Dim) {
+  // Generate the call to get the buffer dimension.
+  switch (CGF->CGM.getTarget().getTriple().getArch()) {
+  case llvm::Triple::dxil:
+    return emitDXILGetDimensions(CGF, Handle, PoisonValue::get(CGF->Int32Ty),
+                                 &Dim);
+    break;
+  case llvm::Triple::spirv:
+    llvm_unreachable("SPIR-V GetDimensions codegen not implemented yet.");
+  default:
+    llvm_unreachable("GetDimensions not supported by target architecture");
+  }
+}
+
+static Value *emitBufferStride(CodeGenFunction *CGF, const Expr *HandleExpr,
+                               LValue &Stride) {
+  // Figure out the stride of the buffer elements from the handle type.
+  auto *HandleTy =
+      cast<HLSLAttributedResourceType>(HandleExpr->getType().getTypePtr());
+  QualType ElementTy = HandleTy->getContainedType();
+  Value *StrideValue = CGF->getTypeSize(ElementTy);
+  return CGF->Builder.CreateStore(StrideValue, Stride.getAddress());
+}
+
 // Return dot product intrinsic that corresponds to the QT scalar type
 static Intrinsic::ID getDotProductIntrinsic(CGHLSLRuntime &RT, QualType QT) {
   if (QT->isFloatingType())
@@ -359,6 +411,15 @@ Value *CodeGenFunction::EmitHLSLBuiltinExpr(unsigned BuiltinID,
         RetTy, CGM.getHLSLRuntime().getNonUniformResourceIndexIntrinsic(),
         ArrayRef<Value *>{IndexOp});
   }
+  case Builtin::BI__builtin_hlsl_buffer_getdimensions: {
+    Value *Handle = EmitScalarExpr(E->getArg(0));
+    LValue Dim = EmitLValue(E->getArg(1));
+    return emitBufferGetDimensions(this, Handle, Dim);
+  }
+  case Builtin::BI__builtin_hlsl_buffer_getstride: {
+    LValue Stride = EmitLValue(E->getArg(1));
+    return emitBufferStride(this, E->getArg(0), Stride);
+  }
   case Builtin::BI__builtin_hlsl_all: {
     Value *Op0 = EmitScalarExpr(E->getArg(0));
     return Builder.CreateIntrinsic(

clang/lib/Sema/HLSLBuiltinTypeDeclBuilder.cpp

@@ -57,6 +57,29 @@ CXXConstructorDecl *lookupCopyConstructor(QualType ResTy) {
       return CD;
   return nullptr;
 }
+
+ParameterABI
+convertParamModifierToParamABI(HLSLParamModifierAttr::Spelling Modifier) {
+  assert(Modifier != HLSLParamModifierAttr::Spelling::Keyword_in &&
+         "HLSL 'in' parameters modifier cannot be converted to ParameterABI");
+  switch (Modifier) {
+  case HLSLParamModifierAttr::Spelling::Keyword_out:
+    return ParameterABI::HLSLOut;
+  case HLSLParamModifierAttr::Spelling::Keyword_inout:
+    return ParameterABI::HLSLInOut;
+  default:
+    llvm_unreachable("Invalid HLSL parameter modifier");
+  }
+}
+
+QualType getInoutParameterType(ASTContext &AST, QualType Ty) {
+  assert(!Ty->isReferenceType() &&
+         "Pointer and reference types cannot be inout or out parameters");
+  Ty = AST.getLValueReferenceType(Ty);
+  Ty.addRestrict();
+  return Ty;
+}
+
 } // namespace
 
 // Builder for template arguments of builtin types. Used internally
@@ -421,13 +444,32 @@ BuiltinTypeMethodBuilder::addParam(StringRef Name, QualType Ty,
 void BuiltinTypeMethodBuilder::createDecl() {
   assert(Method == nullptr && "Method or constructor is already created");
 
-  // create method or constructor type
+  // create function prototype
   ASTContext &AST = DeclBuilder.SemaRef.getASTContext();
   SmallVector<QualType> ParamTypes;
-  for (Param &MP : Params)
-    ParamTypes.emplace_back(MP.Ty);
+  SmallVector<FunctionType::ExtParameterInfo> ParamExtInfos(Params.size());
+  uint32_t ArgIndex = 0;
+  bool IsTemplate = DeclBuilder.Template != nullptr;
+  bool UseParamExtInfo = false;
+  for (Param &MP : Params) {
+    QualType Ty = MP.Ty;
+    if (MP.Modifier != HLSLParamModifierAttr::Keyword_in) {
+      UseParamExtInfo = true;
+      ParamExtInfos[ArgIndex].withABI(
+          convertParamModifierToParamABI(MP.Modifier));
+      // Only update types on inout and out parameters for non-templated
+      // methods. Templated types will have their inout/out parameters
+      // converted to references during template instantiation.
+      if (!IsTemplate)
+        Ty = getInoutParameterType(AST, Ty);
+    }
+    ParamTypes.emplace_back(Ty);
+    ++ArgIndex;
+  }
 
   FunctionProtoType::ExtProtoInfo ExtInfo;
+  if (UseParamExtInfo)
+    ExtInfo.ExtParameterInfos = ParamExtInfos.data();
   if (IsConst)
     ExtInfo.TypeQuals.addConst();
 
@@ -459,6 +501,7 @@ void BuiltinTypeMethodBuilder::createDecl() {
         AST.getTrivialTypeSourceInfo(MP.Ty, SourceLocation()), SC_None,
         nullptr);
     if (MP.Modifier != HLSLParamModifierAttr::Keyword_in) {
+      Parm->setType(getInoutParameterType(AST, Parm->getType()));
       auto *Mod =
           HLSLParamModifierAttr::Create(AST, SourceRange(), MP.Modifier);
       Parm->addAttr(Mod);
@@ -1127,5 +1170,37 @@ BuiltinTypeDeclBuilder &BuiltinTypeDeclBuilder::addConsumeMethod() {
       .finalize();
 }
 
+BuiltinTypeDeclBuilder &
+BuiltinTypeDeclBuilder::addGetDimensionsMethodForBuffer() {
+  using PH = BuiltinTypeMethodBuilder::PlaceHolder;
+  ASTContext &AST = SemaRef.getASTContext();
+  QualType UIntTy = AST.UnsignedIntTy;
+
+  QualType HandleTy = getResourceHandleField()->getType();
+  auto *AttrResTy = cast<HLSLAttributedResourceType>(HandleTy.getTypePtr());
+
+  // Structured buffers except {RW}ByteAddressBuffer have overload
+  // GetDimensions(out uint numStructs, out uint stride).
+  if (AttrResTy->getAttrs().RawBuffer &&
+      AttrResTy->getContainedType() != AST.Char8Ty) {
+    return BuiltinTypeMethodBuilder(*this, "GetDimensions", AST.VoidTy)
+        .addParam("numStructs", UIntTy, HLSLParamModifierAttr::Keyword_out)
+        .addParam("stride", UIntTy, HLSLParamModifierAttr::Keyword_out)
+        .callBuiltin("__builtin_hlsl_buffer_getdimensions", QualType(),
+                     PH::Handle, PH::_0)
+        .callBuiltin("__builtin_hlsl_buffer_getstride", QualType(), PH::Handle,
+                     PH::_1)
+        .finalize();
+  }
+
+  // Typed buffers and {RW}ByteAddressBuffer have overload
+  // GetDimensions(out uint dim).
+  return BuiltinTypeMethodBuilder(*this, "GetDimensions", AST.VoidTy)
+      .addParam("dim", UIntTy, HLSLParamModifierAttr::Keyword_out)
+      .callBuiltin("__builtin_hlsl_buffer_getdimensions", QualType(),
+                   PH::Handle, PH::_0)
+      .finalize();
+}
+
 } // namespace hlsl
 } // namespace clang

clang/lib/Sema/HLSLBuiltinTypeDeclBuilder.h

@@ -95,6 +95,8 @@ class BuiltinTypeDeclBuilder {
   BuiltinTypeDeclBuilder &addAppendMethod();
   BuiltinTypeDeclBuilder &addConsumeMethod();
 
+  BuiltinTypeDeclBuilder &addGetDimensionsMethodForBuffer();
+
 private:
   BuiltinTypeDeclBuilder &addResourceMember(StringRef MemberName,
                                             ResourceClass RC, bool IsROV,

clang/lib/Sema/HLSLExternalSemaSource.cpp

@@ -380,6 +380,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
                     /*RawBuffer=*/false, /*HasCounter=*/false)
         .addArraySubscriptOperators()
         .addLoadMethods()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -392,6 +393,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
                     /*RawBuffer=*/false, /*HasCounter=*/false)
         .addArraySubscriptOperators()
         .addLoadMethods()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -404,6 +406,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
                     /*RawBuffer=*/false, /*HasCounter=*/false)
         .addArraySubscriptOperators()
         .addLoadMethods()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -415,6 +418,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
                     /*RawBuffer=*/true, /*HasCounter=*/false)
         .addArraySubscriptOperators()
         .addLoadMethods()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -428,6 +432,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
         .addLoadMethods()
         .addIncrementCounterMethod()
         .addDecrementCounterMethod()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -439,6 +444,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
     setupBufferType(Decl, *SemaPtr, ResourceClass::UAV, /*IsROV=*/false,
                     /*RawBuffer=*/true, /*HasCounter=*/true)
         .addAppendMethod()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -450,6 +456,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
     setupBufferType(Decl, *SemaPtr, ResourceClass::UAV, /*IsROV=*/false,
                     /*RawBuffer=*/true, /*HasCounter=*/true)
         .addConsumeMethod()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -464,6 +471,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
         .addLoadMethods()
         .addIncrementCounterMethod()
         .addDecrementCounterMethod()
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 
@@ -472,13 +480,15 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
   onCompletion(Decl, [this](CXXRecordDecl *Decl) {
     setupBufferType(Decl, *SemaPtr, ResourceClass::SRV, /*IsROV=*/false,
                     /*RawBuffer=*/true, /*HasCounter=*/false)
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
   Decl = BuiltinTypeDeclBuilder(*SemaPtr, HLSLNamespace, "RWByteAddressBuffer")
              .finalizeForwardDeclaration();
   onCompletion(Decl, [this](CXXRecordDecl *Decl) {
     setupBufferType(Decl, *SemaPtr, ResourceClass::UAV, /*IsROV=*/false,
                     /*RawBuffer=*/true, /*HasCounter=*/false)
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
   Decl = BuiltinTypeDeclBuilder(*SemaPtr, HLSLNamespace,
@@ -487,6 +497,7 @@ void HLSLExternalSemaSource::defineHLSLTypesWithForwardDeclarations() {
   onCompletion(Decl, [this](CXXRecordDecl *Decl) {
     setupBufferType(Decl, *SemaPtr, ResourceClass::UAV, /*IsROV=*/true,
                     /*RawBuffer=*/true, /*HasCounter=*/false)
+        .addGetDimensionsMethodForBuffer()
         .completeDefinition();
   });
 }

clang/lib/Sema/SemaHLSL.cpp

@@ -2978,6 +2978,24 @@ bool SemaHLSL::CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall) {
     TheCall->setType(ResourceTy);
     break;
   }
+  case Builtin::BI__builtin_hlsl_buffer_getdimensions: {
+    ASTContext &AST = SemaRef.getASTContext();
+    if (SemaRef.checkArgCount(TheCall, 2) ||
+        CheckResourceHandle(&SemaRef, TheCall, 0) ||
+        CheckArgTypeMatches(&SemaRef, TheCall->getArg(1), AST.UnsignedIntTy) ||
+        CheckModifiableLValue(&SemaRef, TheCall, 1))
+      return true;
+    break;
+  }
+  case Builtin::BI__builtin_hlsl_buffer_getstride: {
+    ASTContext &AST = SemaRef.getASTContext();
+    if (SemaRef.checkArgCount(TheCall, 2) ||
+        CheckResourceHandle(&SemaRef, TheCall, 0) ||
+        CheckArgTypeMatches(&SemaRef, TheCall->getArg(1), AST.UnsignedIntTy) ||
+        CheckModifiableLValue(&SemaRef, TheCall, 1))
+      return true;
+    break;
+  }
   case Builtin::BI__builtin_hlsl_and:
   case Builtin::BI__builtin_hlsl_or: {
     if (SemaRef.checkArgCount(TheCall, 2))

clang/test/AST/HLSL/ByteAddressBuffers-AST.hlsl

@@ -142,5 +142,19 @@ RESOURCE Buffer;
 // CHECK-NEXT: DeclRefExpr {{.*}} 'hlsl::[[RESOURCE]]' lvalue Var {{.*}} 'tmp' 'hlsl::[[RESOURCE]]'
 // CHECK-NEXT: AlwaysInlineAttr {{.*}} Implicit always_inline
 
+// GetDimensions method
+
+// CHECK-NEXT: CXXMethodDecl {{.*}} GetDimensions 'void (unsigned int &__restrict)'
+// CHECK-NEXT: ParmVarDecl {{.*}} dim 'unsigned int &__restrict'
+// CHECK-NEXT: HLSLParamModifierAttr {{.*}} out
+// CHECK-NEXT: CompoundStmt
+// CHECK-NEXT: CallExpr {{.*}} 'void'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} 'void (*)(...) noexcept' <BuiltinFnToFnPtr>
+// CHECK-NEXT: DeclRefExpr {{.*}} '<builtin fn type>' Function {{.*}} '__builtin_hlsl_buffer_getdimensions' 'void (...) noexcept'
+// CHECK-NEXT: MemberExpr {{.*}} '__hlsl_resource_t {{.*}}' lvalue .__handle {{.*}}
+// CHECK-NEXT: CXXThisExpr {{.*}} 'hlsl::[[RESOURCE]]' lvalue implicit this
+// CHECK-NEXT: DeclRefExpr {{.*}} 'unsigned int' ParmVar {{.*}}  'dim' 'unsigned int &__restrict'
+// CHECK-NEXT: AlwaysInlineAttr {{.*}} Implicit always_inline
+
 // CHECK-NOSUBSCRIPT-NOT: CXXMethodDecl {{.*}} operator[] 'const char8_t &(unsigned int) const'
 // CHECK-NOSUBSCRIPT-NOT: CXXMethodDecl {{.*}} operator[] 'char8_t &(unsigned int)'

clang/test/AST/HLSL/StructuredBuffers-AST.hlsl

@@ -334,6 +334,28 @@ RESOURCE<float> Buffer;
 // CHECK-CONSUME-NEXT: CXXThisExpr {{.*}} 'hlsl::[[RESOURCE]]<element_type>' lvalue implicit this
 // CHECK-CONSUME-NEXT: IntegerLiteral {{.*}} 'int' -1
 
+// GetDimensions method
+
+// CHECK: CXXMethodDecl {{.*}} GetDimensions 'void (unsigned int, unsigned int)'
+// CHECK-NEXT: ParmVarDecl {{.*}} numStructs 'unsigned int &__restrict'
+// CHECK-NEXT: HLSLParamModifierAttr {{.*}} out
+// CHECK-NEXT: ParmVarDecl {{.*}} stride 'unsigned int &__restrict'
+// CHECK-NEXT: HLSLParamModifierAttr {{.*}} out
+// CHECK-NEXT: CompoundStmt
+// CHECK-NEXT: CallExpr {{.*}} 'void'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} 'void (*)(...) noexcept' <BuiltinFnToFnPtr>
+// CHECK-NEXT: DeclRefExpr {{.*}} '<builtin fn type>' Function {{.*}} '__builtin_hlsl_buffer_getdimensions' 'void (...) noexcept'
+// CHECK-NEXT: MemberExpr {{.*}} '__hlsl_resource_t {{.*}}' lvalue .__handle {{.*}}
+// CHECK-NEXT: CXXThisExpr {{.*}} 'hlsl::[[RESOURCE]]<element_type>' lvalue implicit this
+// CHECK-NEXT: DeclRefExpr {{.*}} 'unsigned int' ParmVar {{.*}}  'numStructs' 'unsigned int &__restrict'
+// CHECK-NEXT: CallExpr {{.*}} 'void'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} 'void (*)(...) noexcept' <BuiltinFnToFnPtr>
+// CHECK-NEXT: DeclRefExpr {{.*}} '<builtin fn type>' Function {{.*}} '__builtin_hlsl_buffer_getstride' 'void (...) noexcept'
+// CHECK-NEXT: MemberExpr {{.*}} '__hlsl_resource_t {{.*}}' lvalue .__handle {{.*}}
+// CHECK-NEXT: CXXThisExpr {{.*}} 'hlsl::[[RESOURCE]]<element_type>' lvalue implicit this
+// CHECK-NEXT: DeclRefExpr {{.*}} 'unsigned int' ParmVar {{.*}}  'stride' 'unsigned int &__restrict'
+// CHECK-NEXT: AlwaysInlineAttr {{.*}} Implicit always_inline
+
 // CHECK: ClassTemplateSpecializationDecl {{.*}} class [[RESOURCE]] definition
 // CHECK: TemplateArgument type 'float'
 // CHECK-NEXT: BuiltinType {{.*}} 'float'

clang/test/AST/HLSL/TypedBuffers-AST.hlsl

@@ -214,6 +214,20 @@ RESOURCE<float> Buffer;
 // CHECK-NEXT: DeclRefExpr {{.*}} 'unsigned int' ParmVar {{.*}}  'Index' 'unsigned int'
 // CHECK-NEXT: AlwaysInlineAttr {{.*}} Implicit always_inline
 
+// GetDimensions method
+
+// CHECK-NEXT: CXXMethodDecl {{.*}} GetDimensions 'void (unsigned int)'
+// CHECK-NEXT: ParmVarDecl {{.*}} dim 'unsigned int &__restrict'
+// CHECK-NEXT: HLSLParamModifierAttr {{.*}} out
+// CHECK-NEXT: CompoundStmt
+// CHECK-NEXT: CallExpr {{.*}} 'void'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} 'void (*)(...) noexcept' <BuiltinFnToFnPtr>
+// CHECK-NEXT: DeclRefExpr {{.*}} '<builtin fn type>' Function {{.*}} '__builtin_hlsl_buffer_getdimensions' 'void (...) noexcept'
+// CHECK-NEXT: MemberExpr {{.*}} '__hlsl_resource_t {{.*}}' lvalue .__handle {{.*}}
+// CHECK-NEXT: CXXThisExpr {{.*}} 'hlsl::[[RESOURCE]]<element_type>' lvalue implicit this
+// CHECK-NEXT: DeclRefExpr {{.*}} 'unsigned int' ParmVar {{.*}}  'dim' 'unsigned int &__restrict'
+// CHECK-NEXT: AlwaysInlineAttr {{.*}} Implicit always_inline
+
 // CHECK: ClassTemplateSpecializationDecl {{.*}} class [[RESOURCE]] definition
 
 // CHECK: TemplateArgument type 'float'