[HLSL] Generate buffer subscript operators

[beanz]

Summary

In HLSL buffer types support array subscripting syntax for loads and
stores. This change fleshes out the subscript operators to become array
accesses on the underlying handle pointer. This will allow LLVM
optimization passes to optimize resource accesses the same way any other
memory access would be optimized.

[python3kgae]

clang/lib/Sema/HLSLExternalSemaSource.cpp

221

If we reuse this function for StructuredBuffer, then const subscript return a const reference could be better?

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

5

Maybe change In to RWBuffer<float4>
And Out[Idx] = In[Idx].z;
So we also test vector case?

17

Where is the this.addr coming from?

[beanz]

clang/lib/Sema/HLSLExternalSemaSource.cpp

221

Why? Since this should get inlined return value optimization should eliminate the redundant copies.

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

5

None of the added code here depends on the template parameter type. Putting a vector in increases the complexity of the output, but I'm unsure it extends the test coverage meaningfully.

17

That is how clang generates stored locations for input parameters.

[python3kgae]

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

4

Why add const instead of using Buffer directly?

[beanz]

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

4

Making this const forces the const methods to be used. It is just to drive the correct validation and code generation.

[python3kgae]

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

4

So maybe we don't need Buffer at all, just use const RWBuffer for read-only usage?

[beanz]

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

4

I think there would likely be some code incompatibilities if we alias Buffer to const RWBuffer, but that would be interesting to consider.

[beanz]

Updating based on PR feedback and rebasing.

• Rebased on main today
• Made const subscript return type const &

[aaron.ballman]

clang/lib/Sema/HLSLExternalSemaSource.cpp

108–113

180

Can you use a static_cast to make this a wee bit safer, or is that not possible?

232–233

242

244

Should this be using a size type instead?

253

255

258

263

267

[beanz]

clang/lib/Sema/HLSLExternalSemaSource.cpp

244

As currently implemented it is an unsigned int.

I should include the doc link in the description:
https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sm5-object-rwbuffer-operatorindex

Eventually I think we'll need to introduce size types to HLSL, but we currently don't really use them.

[beanz]

Updating based on @aaron.ballman's feedback.

• Change reinterpret_cast -> static_cast
• Aaron likes auto more than me... but all in good places :D

[aaron.ballman]

LGTM aside from a nit that was missed. Thanks for switching to static_cast, that makes me happier. :-)

clang/lib/Sema/HLSLExternalSemaSource.cpp

109

Missed from the previous suggestion.

clang/lib/Sema/HLSLExternalSemaSource.cpp

addMemberVariable(StringRef Name, QualType Type,
     Field->setImplicit(true);
     Record->addDecl(Field);
     Fields[Name] = Field;
     return *this;
   }
 
   BuiltinTypeDeclBuilder &
   addHandleMember(AccessSpecifier Access = AccessSpecifier::AS_private) {
-    return addMemberVariable("h", Record->getASTContext().VoidPtrTy, Access);
+    QualType Ty = Record->getASTContext().VoidPtrTy;
+    if (Template) {
+      if (const auto *TTD = dyn_cast<TemplateTypeParmDecl>(
+              Template->getTemplateParameters()->getParam(0)))
+        Ty = Record->getASTContext().getPointerType(
+            QualType(TTD->getTypeForDecl(), 0));
+    }
+    return addMemberVariable("h", Ty, Access);
   }
 
   BuiltinTypeDeclBuilder &
   annotateResourceClass(HLSLResourceAttr::ResourceClass RC) {
     Record->addAttr(
         HLSLResourceAttr::CreateImplicit(Record->getASTContext(), RC));
     return *this;
   }
CXXConstructorDecl *Constructor = CXXConstructorDecl::Create(
         AST.getTrivialTypeSourceInfo(ConstructorType, SourceLocation()),
         ExplicitSpecifier(), false, true, false,
         ConstexprSpecKind::Unspecified);
 
     DeclRefExpr *Fn =
         lookupBuiltinFunction(AST, S, "__builtin_hlsl_create_handle");
 
     Expr *RCExpr = emitResourceClassExpr(AST, RC);
-    CallExpr *Call =
-        CallExpr::Create(AST, Fn, {RCExpr}, AST.VoidPtrTy, VK_PRValue,
+    Expr *Call = CallExpr::Create(AST, Fn, {RCExpr}, AST.VoidPtrTy, VK_PRValue,
                                   SourceLocation(), FPOptionsOverride());
 
     CXXThisExpr *This = new (AST)
         CXXThisExpr(SourceLocation(), Constructor->getThisType(), true);
-    MemberExpr *Handle = MemberExpr::CreateImplicit(
-        AST, This, true, Fields["h"], Fields["h"]->getType(), VK_LValue,
+    Expr *Handle = MemberExpr::CreateImplicit(AST, This, true, Fields["h"],
+                                              Fields["h"]->getType(), VK_LValue,
                                               OK_Ordinary);
+
+    // If the handle isn't a void pointer, cast the builtin result to the
+    // correct type.
+    if (Handle->getType().getCanonicalType() != AST.VoidPtrTy) {
+      Call = CXXStaticCastExpr::Create(
+          AST, Handle->getType(), VK_PRValue, CK_Dependent, Call, nullptr,
+          AST.getTrivialTypeSourceInfo(Handle->getType(), SourceLocation()),
+          FPOptionsOverride(), SourceLocation(), SourceLocation(),
+          SourceRange());
+    }
+
     BinaryOperator *Assign = BinaryOperator::Create(
         AST, Handle, Call, BO_Assign, Handle->getType(), VK_LValue, OK_Ordinary,
         SourceLocation(), FPOptionsOverride());
 
     Constructor->setBody(
         CompoundStmt::Create(AST, {Assign}, FPOptionsOverride(),
                              SourceLocation(), SourceLocation()));
     Constructor->setAccess(AccessSpecifier::AS_public);
     Record->addDecl(Constructor);
     return *this;
   }
 
+  BuiltinTypeDeclBuilder &addArraySubscriptOperators() {
+    addArraySubscriptOperator(true);
+    addArraySubscriptOperator(false);
+    return *this;
+  }
+
+  BuiltinTypeDeclBuilder &addArraySubscriptOperator(bool IsConst) {
+    assert(Fields.count("h") > 0 &&
+           "Subscript operator must be added after the handle.");
+
+    FieldDecl *Handle = Fields["h"];
+    ASTContext &AST = Record->getASTContext();
+
+    assert(Handle->getType().getCanonicalType() != AST.VoidPtrTy &&
+           "Not yet supported for void pointer handles.");
+
+    QualType ElemTy =
+        QualType(Handle->getType()->getPointeeOrArrayElementType(), 0);
+    QualType ReturnTy = ElemTy;
+
+    FunctionProtoType::ExtProtoInfo ExtInfo;
+
+    // Subscript operators return references to elements, const makes the
+    // reference and method const so that the underlying data is not mutable.
+    ReturnTy = AST.getLValueReferenceType(ReturnTy);
+    if (IsConst) {
+      ExtInfo.TypeQuals.addConst();
+      ReturnTy.addConst();
+    }
+
+    QualType MethodTy =
+        AST.getFunctionType(ReturnTy, {AST.UnsignedIntTy}, ExtInfo);
+    auto *TSInfo = AST.getTrivialTypeSourceInfo(MethodTy, SourceLocation());
+    auto *MethodDecl = CXXMethodDecl::Create(
+        AST, Record, SourceLocation(),
+        DeclarationNameInfo(
+            AST.DeclarationNames.getCXXOperatorName(OO_Subscript),
+            SourceLocation()),
+        MethodTy, TSInfo, SC_None, false, false, ConstexprSpecKind::Unspecified,
+        SourceLocation());
+
+    IdentifierInfo &II = AST.Idents.get("Idx", tok::TokenKind::identifier);
+    auto *IdxParam = ParmVarDecl::Create(
+        AST, MethodDecl->getDeclContext(), SourceLocation(), SourceLocation(),
+        &II, AST.UnsignedIntTy,
+        AST.getTrivialTypeSourceInfo(AST.UnsignedIntTy, SourceLocation()),
+        SC_None, nullptr);
+    MethodDecl->setParams({IdxParam});
+
+    // Also add the parameter to the function prototype.
+    auto FnProtoLoc = TSInfo->getTypeLoc().getAs<FunctionProtoTypeLoc>();
+    FnProtoLoc.setParam(0, IdxParam);
+
+    auto *This = new (AST)
+        CXXThisExpr(SourceLocation(), MethodDecl->getThisType(), true);
+    auto *HandleAccess = MemberExpr::CreateImplicit(
+        AST, This, true, Handle, Handle->getType(), VK_LValue, OK_Ordinary);
+
+    auto *IndexExpr = DeclRefExpr::Create(
+        AST, NestedNameSpecifierLoc(), SourceLocation(), IdxParam, false,
+        DeclarationNameInfo(IdxParam->getDeclName(), SourceLocation()),
+        AST.UnsignedIntTy, VK_PRValue);
+
+    auto *Array =
+        new (AST) ArraySubscriptExpr(HandleAccess, IndexExpr, ElemTy, VK_LValue,
+                                     OK_Ordinary, SourceLocation());
+
+    auto *Return = ReturnStmt::Create(AST, SourceLocation(), Array, nullptr);
+
+    MethodDecl->setBody(CompoundStmt::Create(AST, {Return}, FPOptionsOverride(),
+                                             SourceLocation(),
+                                             SourceLocation()));
+    MethodDecl->setLexicalDeclContext(Record);
+    MethodDecl->setAccess(AccessSpecifier::AS_public);
+    Record->addDecl(MethodDecl);
+
+    return *this;
+  }
+
   BuiltinTypeDeclBuilder &startDefinition() {
     Record->startDefinition();
     return *this;
   }
 
   BuiltinTypeDeclBuilder &completeDefinition() {
     assert(Record->isBeingDefined() &&
            "Definition must be started before completing it.");
if (It == Completions.end())
     return;
   It->second(Record);
 }
 
 void HLSLExternalSemaSource::completeBufferType(CXXRecordDecl *Record) {
   BuiltinTypeDeclBuilder(Record)
       .addHandleMember()
       .addDefaultHandleConstructor(*SemaPtr, ResourceClass::UAV)
+      .addArraySubscriptOperators()
       .annotateResourceClass(HLSLResourceAttr::UAV)
       .completeDefinition();
 }

clang/lib/Sema/SemaType.cpp

QualType Sema::BuildPointerType(QualType T,
 
   if (T->isFunctionType() && getLangOpts().OpenCL &&
       !getOpenCLOptions().isAvailableOption("__cl_clang_function_pointers",
                                             getLangOpts())) {
     Diag(Loc, diag::err_opencl_function_pointer) << /*pointer*/ 0;
     return QualType();
   }
 
-  if (getLangOpts().HLSL) {
+  if (getLangOpts().HLSL && Loc.isValid()) {
     Diag(Loc, diag::err_hlsl_pointers_unsupported) << 0;
     return QualType();
   }
 
   if (checkQualifiedFunction(*this, T, Loc, QFK_Pointer))
     return QualType();
 
   assert(!T->isObjCObjectType() && "Should build ObjCObjectPointerType");
QualType Sema::BuildReferenceType(QualType T, bool SpelledAsLValue,
   // C++ [dcl.ref]p1:
   //   A declarator that specifies the type "reference to cv void"
   //   is ill-formed.
   if (T->isVoidType()) {
     Diag(Loc, diag::err_reference_to_void);
     return QualType();
   }
 
-  if (getLangOpts().HLSL) {
+  if (getLangOpts().HLSL && Loc.isValid()) {
     Diag(Loc, diag::err_hlsl_pointers_unsupported) << 1;
     return QualType();
   }
 
   if (checkQualifiedFunction(*this, T, Loc, QFK_Reference))
     return QualType();
 
   if (T->isFunctionType() && getLangOpts().OpenCL &&
QualType Sema::BuildMemberPointerType(QualType T, QualType Class,
 
   if (T->isFunctionType() && getLangOpts().OpenCL &&
       !getOpenCLOptions().isAvailableOption("__cl_clang_function_pointers",
                                             getLangOpts())) {
     Diag(Loc, diag::err_opencl_function_pointer) << /*pointer*/ 0;
     return QualType();
   }
 
-  if (getLangOpts().HLSL) {
+  if (getLangOpts().HLSL && Loc.isValid()) {
     Diag(Loc, diag::err_hlsl_pointers_unsupported) << 0;
     return QualType();
   }
 
   // Adjust the default free function calling convention to the default method
   // calling convention.
   bool IsCtorOrDtor =
       (Entity.getNameKind() == DeclarationName::CXXConstructorName) ||

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

 // CHECK: ClassTemplateDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit RWBuffer
 // CHECK-NEXT: TemplateTypeParmDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> class depth 0 index 0 element_type
 // CHECK-NEXT: TemplateArgument type 'float'
 // CHECK-NEXT: BuiltinType 0x{{[0-9A-Fa-f]+}} 'float'
 // CHECK-NEXT: CXXRecordDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit class RWBuffer definition
 
 // CHECK: FinalAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit final
 // CHECK-NEXT: HLSLResourceAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit UAV
-// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit h 'void *'
+// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> implicit h 'element_type *'
+
+// CHECK: CXXMethodDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> operator[] 'element_type &const (unsigned int) const'
+// CHECK-NEXT: ParmVarDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> Idx 'unsigned int'
+// CHECK-NEXT: CompoundStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
+// CHECK-NEXT: ReturnStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
+// CHECK-NEXT: ArraySubscriptExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type' lvalue
+// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type *' lvalue ->h 0x{{[0-9A-Fa-f]+}}
+// CHECK-NEXT: CXXThisExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'const RWBuffer<element_type> *' implicit this
+// CHECK-NEXT: DeclRefExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'unsigned int' ParmVar 0x{{[0-9A-Fa-f]+}} 'Idx' 'unsigned int'
+
+// CHECK-NEXT: CXXMethodDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> operator[] 'element_type &(unsigned int)'
+// CHECK-NEXT: ParmVarDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> Idx 'unsigned int'
+// CHECK-NEXT: CompoundStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
+// CHECK-NEXT: ReturnStmt 0x{{[0-9A-Fa-f]+}} <<invalid sloc>>
+// CHECK-NEXT: ArraySubscriptExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type' lvalue
+// CHECK-NEXT: MemberExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'element_type *' lvalue ->h 0x{{[0-9A-Fa-f]+}}
+// CHECK-NEXT: CXXThisExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'RWBuffer<element_type> *' implicit this
+// CHECK-NEXT: DeclRefExpr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> 'unsigned int' ParmVar 0x{{[0-9A-Fa-f]+}} 'Idx' 'unsigned int'
+
 // CHECK: ClassTemplateSpecializationDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc> class RWBuffer definition
 
 // CHECK: TemplateArgument type 'float'
 // CHECK-NEXT: BuiltinType 0x{{[0-9A-Fa-f]+}} 'float'
 // CHECK-NEXT: FinalAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit final
 // CHECK-NEXT: HLSLResourceAttr 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> Implicit UAV
-// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc>  implicit referenced h 'void *'
+// CHECK-NEXT: FieldDecl 0x{{[0-9A-Fa-f]+}} <<invalid sloc>> <invalid sloc>  implicit referenced h 'float *'

clang/test/CodeGenHLSL/buffer-array-operator.hlsl

+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -emit-llvm -disable-llvm-passes -o - %s | FileCheck %s
+
+const RWBuffer<float> In;
+RWBuffer<float> Out;
+
+void fn(int Idx) {
+  Out[Idx] = In[Idx];
+}
+
+// This test is intended to verify reasonable code generation of the subscript
+// operator. In this test case we should be generating both the const and
+// non-const operators so we verify both cases.
+
+// Non-const comes first.
+// CHECK: ptr @"??A?$RWBuffer@M@hlsl@@QBAAAMI@Z"
+// CHECK: %this1 = load ptr, ptr %this.addr, align 4
+// CHECK-NEXT: %h = getelementptr inbounds %"class.hlsl::RWBuffer", ptr %this1, i32 0, i32 0
+// CHECK-NEXT: %0 = load ptr, ptr %h, align 4
+// CHECK-NEXT: %1 = load i32, ptr %Idx.addr, align 4
+// CHECK-NEXT: %arrayidx = getelementptr inbounds float, ptr %0, i32 %1
+// CHECK-NEXT: ret ptr %arrayidx
+
+// Const comes next, and returns the pointer instead of the value.
+// CHECK: ptr @"??A?$RWBuffer@M@hlsl@@QAAAAMI@Z"
+// CHECK: %this1 = load ptr, ptr %this.addr, align 4
+// CHECK-NEXT: %h = getelementptr inbounds %"class.hlsl::RWBuffer", ptr %this1, i32 0, i32 0
+// CHECK-NEXT: %0 = load ptr, ptr %h, align 4
+// CHECK-NEXT: %1 = load i32, ptr %Idx.addr, align 4
+// CHECK-NEXT: %arrayidx = getelementptr inbounds float, ptr %0, i32 %1
+// CHECK-NEXT: ret ptr %arrayidx