[HLSL] Support PCH for cc1 mode

[python3kgae]

Summary

Add HLSLExternalSemaSource as ExternalSemaSource instead of ASTContext::ExternalSource when PCH is included.

This allows a different external source to be set for the AST context.

[beanz]

This change doesn't look like it is NFC to me. If I'm understanding correctly it results in the HLSLExternalSemaSource being set for Sema, but allows a different external source to be set for the AST context. That doesn't seem unreasonable to me, but perhaps we should have a test case that verifies that it works correctly?

Perhaps a test case with a PCH that uses something populated by the HLSL external source? That would demonstrate that the AST source for the PCH works and that it doesn't interfere with the HLSL source.

[python3kgae]

When there's no external source set, add HLSLSema as ExternalSource so ASTWriter can access HLSLSema.

[python3kgae]

Add newline at end of file

[beanz]

clang/test/SemaHLSL/pch.hlsl

5 ↗

(On Diff #454889)

I want to make sure I'm following this test.

You compile this file to a PCH (%t), then compile /dev/null using this %t as the PCH.

IIUC, that doesn't actually verify that PCH and the HLSL sema work together. There is no test here that pulls from the PCH and HLSL Sema in the same translation unit.

My suggestion would be to:
(1) move this test under AST/HLSL (since it is an AST test, not a Sema test)
(2) put a header under AST/HLSL/Inputs that (like this file) contains a function that uses HLSL types
(3) make the test file separate from the header, and have it use both the function in the header and some completely different HLSL-provided type (like declare a RWBuffer).
(4) your test can then generate the PCH, and use the PCH to compile the source verifying that both AST sources work

[python3kgae]

Reuse decls already in PCH when initailize HLSLExternalSemaSource.

[python3kgae]

Fix build error

[python3kgae]

Fix format.

[beanz]

clang/include/clang/Sema/HLSLExternalSemaSource.h

57

IIUC, this code just exists to make sure that the ASTReader deserializes before the external sema source starts adding things. Is that correct?

If so, you don't need to do this, instead you can just add this code to InitializeSema() to force the ASTReader to de-serialize the decls:

// If the translation unit has external storage force external decls to load.
if (AST.getTranslationUnitDecl()->hasExternalLexicalStorage())
    (void)AST.getTranslationUnitDecl()->decls_begin();

clang/lib/Sema/HLSLExternalSemaSource.cpp

399

I think the core of what you'e doing here is not far off, but I think this pattern doesn't scale or cover all the use cases that matter.

The cases that come to my mind are:

1. PCH has a forward declaration, source doesn't use the type so we shouldn't complete the type.
2. PCH has a forward declaration, source uses the type, so we need to generate a complete decl.
3. PCH has a complete declaration, so we should re-use the declaration and not need to do anything.

We also need to keep in mind that we're going to keep adding more and more types, so this pattern where each new type needs to add code to lookup is going to be a challenge. It is also going to be a challenge to manage the complexity of PCH defines one thing but not another.

One thing you should look at is how clang handles the following C++ code:

template<typename T>
class Foo;

template<typename T>
class Foo {
  T Val;
};

This results in two decls. One for the forward decl and the second for the definition. The definition lists the forward declaration as the previous decl. We should do that here.

When we generate the HLSL namespace, we can do a lookup and if we find a previous decl, set the previous decl.

When we generate any of the builtin types, we can lookup the previous decl, and annotate as the previous decl. We'll can also handle the case where a decl from the PCH is complete. If the PCH provides a complete decl, we can skip generating any decls for it, and the lookups should just work.

[python3kgae]

clang/lib/Sema/HLSLExternalSemaSource.cpp

399

Set previous decl not work.
It goes to "For most kinds of declaration, it doesn't really matter which one we pick." before checking PreviousDecl.

   // For most kinds of declaration, it doesn't really matter which one we pick.
  if (!isa<FunctionDecl>(DUnderlying) && !isa<VarDecl>(DUnderlying)) {
    // If the existing declaration is hidden, prefer the new one. Otherwise,
    // keep what we've got.
    return !S.isVisible(Existing);
  }



  // Pick the newer declaration; it might have a more precise type.
  for (Decl *Prev = DUnderlying->getPreviousDecl(); Prev;
       Prev = Prev->getPreviousDecl())
    if (Prev == EUnderlying)
      return true;
  return false;

[python3kgae]

Move reuse decl in PCH into BuiltinTypeDeclBuilder.

[python3kgae]

clang/include/clang/Sema/HLSLExternalSemaSource.h

57

Still need this to make sure HLSLSema and ExternalSema are initialized before MultiplexExternalSemaSource.

[beanz]

clang/include/clang/Sema/HLSLExternalSemaSource.h

57

In FrontendAction, where you are creating the Chained source, you can instead create a Multiplex source, and put the PCH external source in as the first source and the HLSL one second. The PCH will get initialized first, the HLSL one can force the PCH one to populate the decls.

clang/lib/Sema/HLSLExternalSemaSource.cpp

397

Rather than reusing the namespace you can set the PCH one as the previous decl for the one you're declaring here. That will result in them basically being merged because Namespaces can have multiple decls.

The advantage to this is that the PCH namespace will be marked in the AST as a deserialized decl but the HLSL one won't. It will make debugging and AST visualization easier.

479

If the decl you put in here is the canonical decl, and you do the lookup by canonical decl. The conflicting decls shouldn't be an issue. That does require that you mark the previous decl so that they are known to be the same declaration.

[python3kgae]

clang/include/clang/Sema/HLSLExternalSemaSource.h

57

Asked how to resolve the issue in https://discourse.llvm.org/t/is-it-possible-for-multiplexexternalsemasource-to-own-the-sources/64990

Hope someone can explain why multiplexexternalsemasource doesn't own the Sources.

[python3kgae]

Rebase to use MultiplexExternalSemaSource.

[python3kgae]

clang/lib/Sema/HLSLExternalSemaSource.cpp

62

If comment this line out. The test will pass.

[python3kgae]

Merge fix from Chris.

[python3kgae]

Only add HLSLSema when not has it in ExternalSource.

[python3kgae]

Only add HLSLExternalSemaSource when ExternalSource is not HLSLExternalSemaSource.

[beanz]

clang/lib/Frontend/FrontendAction.cpp

1036

This seems off. The multi source should always have at least two external sema sources.

Is there any situation where the HLSL sema would already be added?

[beanz]

Looping @aaron.ballman in here too. I think we need some eyes from outside our team on this before landing it.

[python3kgae]

Remove useless code since ExternalSource is not be reused when there're more than one Input files.

[python3kgae]

clang/lib/Frontend/FrontendAction.cpp

1036

Remove. I thought BeginSourceFile would reuse the same ExternalSource.

[python3kgae]

Fix test fail caused by not add HLSL builtin Resource type.

[aaron.ballman]

clang/lib/Sema/HLSLExternalSemaSource.cpp

50

99.9% sure this is safe because we're looking up a tag name and I can't think of any way that might return multiple results for the same identifier, but bringing it up just in case anyone else knows of some obscure extension where that's possible.

55

Is it possible that this finds a different kind of tag, like an enumeration?

85–86

A downside to this pattern is that we need to repeat the "if it's a complete definition, don't do anything" predicate in basically any mutating member function. Getting that wrong will lead to hard-to-debug issues with PCH, as I understand it. But I don't think there's a cleaner way to do that without doing far more invasive changes.

378

How certain are you that there's only one result possible here?

387

459–461

We made a lookup result but then do nothing with it?

[python3kgae]

Code cleanup per comments.

clang/lib/Sema/HLSLExternalSemaSource.cpp

55

No, enum with the same name will get an error when building the PCH.

378

I cannot think of a case more than one result is here.

459–461

Removed.

[python3kgae]

Rebase

[beanz]

One small nit, otherwise looks good.

clang/include/clang/Sema/HLSLExternalSemaSource.h

18

nit: This seems to be unused

[python3kgae]

Move #include "clang/Sema/MultiplexExternalSemaSource.h" into FrontendAction.cpp

clang/include/clang/Sema/HLSLExternalSemaSource.h

 //
 //===----------------------------------------------------------------------===//
 #ifndef CLANG_SEMA_HLSLEXTERNALSEMASOURCE_H
 #define CLANG_SEMA_HLSLEXTERNALSEMASOURCE_H
 
 #include "llvm/ADT/DenseMap.h"
 
 #include "clang/Sema/ExternalSemaSource.h"
+#include "clang/Sema/MultiplexExternalSemaSource.h"
 
 namespace clang {
 class NamespaceDecl;
 class Sema;
 
 class HLSLExternalSemaSource : public ExternalSemaSource {
   Sema *SemaPtr = nullptr;
-  NamespaceDecl *HLSLNamespace;
+  NamespaceDecl *HLSLNamespace = nullptr;
   CXXRecordDecl *ResourceDecl;
 
   using CompletionFunction = std::function<void(CXXRecordDecl *)>;
   llvm::DenseMap<CXXRecordDecl *, CompletionFunction> Completions;
 
   void defineHLSLVectorAlias();
   void defineTrivialHLSLTypes();
   void forwardDeclareHLSLTypes();
public:
 
   using ExternalASTSource::CompleteType;
   /// Complete an incomplete HLSL builtin type
   void CompleteType(TagDecl *Tag) override;
 };
 
 } // namespace clang
 
 #endif // CLANG_SEMA_HLSLEXTERNALSEMASOURCE_H

clang/lib/Frontend/FrontendAction.cpp

if (!CI.getFrontendOpts().OverrideRecordLayoutsFile.empty() &&
     IntrusiveRefCntPtr<ExternalASTSource>
       Override(new LayoutOverrideSource(
                      CI.getFrontendOpts().OverrideRecordLayoutsFile));
     CI.getASTContext().setExternalSource(Override);
   }
 
   // Setup HLSL External Sema Source
   if (CI.getLangOpts().HLSL && CI.hasASTContext()) {
-    IntrusiveRefCntPtr<ExternalASTSource> HLSLSema(
+    IntrusiveRefCntPtr<ExternalSemaSource> HLSLSema(
         new HLSLExternalSemaSource());
+    if (auto *SemaSource = dyn_cast_if_present<ExternalSemaSource>(
+            CI.getASTContext().getExternalSource())) {
+      IntrusiveRefCntPtr<ExternalSemaSource> MultiSema(
+          new MultiplexExternalSemaSource(SemaSource, HLSLSema.get()));
+      CI.getASTContext().setExternalSource(MultiSema);
+    } else
       CI.getASTContext().setExternalSource(HLSLSema);
   }
 
   FailureCleanup.release();
   return true;
 }
 
 llvm::Error FrontendAction::Execute() {
   CompilerInstance &CI = getCompilerInstance();

clang/lib/Sema/HLSLExternalSemaSource.cpp

 
 namespace {
 
 struct TemplateParameterListBuilder;
 
 struct BuiltinTypeDeclBuilder {
   CXXRecordDecl *Record = nullptr;
   ClassTemplateDecl *Template = nullptr;
+  ClassTemplateDecl *PrevTemplate = nullptr;
   NamespaceDecl *HLSLNamespace = nullptr;
   llvm::StringMap<FieldDecl *> Fields;
 
   BuiltinTypeDeclBuilder(CXXRecordDecl *R) : Record(R) {
     Record->startDefinition();
     Template = Record->getDescribedClassTemplate();
   }
 
   BuiltinTypeDeclBuilder(Sema &S, NamespaceDecl *Namespace, StringRef Name)
       : HLSLNamespace(Namespace) {
     ASTContext &AST = S.getASTContext();
     IdentifierInfo &II = AST.Idents.get(Name, tok::TokenKind::identifier);
 
+    LookupResult Result(S, &II, SourceLocation(), Sema::LookupTagName);
+    CXXRecordDecl *PrevDecl = nullptr;
+    if (S.LookupQualifiedName(Result, HLSLNamespace)) {
+      NamedDecl *Found = Result.getFoundDecl();
+      if (auto *TD = dyn_cast<ClassTemplateDecl>(Found)) {
+        PrevDecl = TD->getTemplatedDecl();
+        PrevTemplate = TD;
+      } else
+        PrevDecl = dyn_cast<CXXRecordDecl>(Found);
+      assert(PrevDecl && "Unexpected lookup result type.");
+    }
+
+    if (PrevDecl && PrevDecl->isCompleteDefinition()) {
+      Record = PrevDecl;
+      return;
+    }
+
     Record = CXXRecordDecl::Create(AST, TagDecl::TagKind::TTK_Class,
                                    HLSLNamespace, SourceLocation(),
-                                   SourceLocation(), &II, nullptr, true);
+                                   SourceLocation(), &II, PrevDecl, true);
     Record->setImplicit(true);
     Record->setLexicalDeclContext(HLSLNamespace);
     Record->setHasExternalLexicalStorage();
 
-    // Don't let anyone derive from built-in types
+    // Don't let anyone derive from built-in types.
     Record->addAttr(FinalAttr::CreateImplicit(AST, SourceRange(),
                                               AttributeCommonInfo::AS_Keyword,
                                               FinalAttr::Keyword_final));
   }
 
   ~BuiltinTypeDeclBuilder() {
-    if (HLSLNamespace && !Template)
+    if (!Record->getLexicalDeclContext())
       HLSLNamespace->addDecl(Record);
   }
 
   BuiltinTypeDeclBuilder &
-  addTemplateArgumentList(llvm::ArrayRef<NamedDecl *> TemplateArgs) {
-    ASTContext &AST = Record->getASTContext();
-
-    auto *ParamList =
-        TemplateParameterList::Create(AST, SourceLocation(), SourceLocation(),
-                                      TemplateArgs, SourceLocation(), nullptr);
-    Template = ClassTemplateDecl::Create(
-        AST, Record->getDeclContext(), SourceLocation(),
-        DeclarationName(Record->getIdentifier()), ParamList, Record);
-    Record->setDescribedClassTemplate(Template);
-    Template->setImplicit(true);
-    Template->setLexicalDeclContext(Record->getDeclContext());
-    Record->getDeclContext()->addDecl(Template);
-
-    // Requesting the class name specialization will fault in required types.
-    QualType T = Template->getInjectedClassNameSpecialization();
-    T = AST.getInjectedClassNameType(Record, T);
-    return *this;
-  }
-
-  BuiltinTypeDeclBuilder &
   addMemberVariable(StringRef Name, QualType Type,
                     AccessSpecifier Access = AccessSpecifier::AS_private) {
+    if (Record->isCompleteDefinition())
+      return *this;
     assert(Record->isBeingDefined() &&
            "Definition must be started before adding members!");
     ASTContext &AST = Record->getASTContext();
 
     IdentifierInfo &II = AST.Idents.get(Name, tok::TokenKind::identifier);
     TypeSourceInfo *MemTySource =
         AST.getTrivialTypeSourceInfo(Type, SourceLocation());
     auto *Field = FieldDecl::Create(
         AST, Record, SourceLocation(), SourceLocation(), &II, Type, MemTySource,
         nullptr, false, InClassInitStyle::ICIS_NoInit);
     Field->setAccess(Access);
     Field->setImplicit(true);
     Record->addDecl(Field);
     Fields[Name] = Field;
     return *this;
   }
 
   BuiltinTypeDeclBuilder &
   addHandleMember(AccessSpecifier Access = AccessSpecifier::AS_private) {
+    if (Record->isCompleteDefinition())
+      return *this;
     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) {
+    if (Record->isCompleteDefinition())
+      return *this;
     Record->addAttr(
         HLSLResourceAttr::CreateImplicit(Record->getASTContext(), RC));
     return *this;
   }
 
   static DeclRefExpr *lookupBuiltinFunction(ASTContext &AST, Sema &S,
                                             StringRef Name) {
     CXXScopeSpec SS;
return IntegerLiteral::Create(
         AST,
         llvm::APInt(AST.getIntWidth(AST.UnsignedCharTy),
                     static_cast<uint8_t>(RC)),
         AST.UnsignedCharTy, SourceLocation());
   }
 
   BuiltinTypeDeclBuilder &addDefaultHandleConstructor(Sema &S,
                                                       ResourceClass RC) {
+    if (Record->isCompleteDefinition())
+      return *this;
     ASTContext &AST = Record->getASTContext();
 
     QualType ConstructorType =
         AST.getFunctionType(AST.VoidTy, {}, FunctionProtoType::ExtProtoInfo());
 
     CanQualType CanTy = Record->getTypeForDecl()->getCanonicalTypeUnqualified();
     DeclarationName Name = AST.DeclarationNames.getCXXConstructorName(CanTy);
     CXXConstructorDecl *Constructor = CXXConstructorDecl::Create(
Constructor->setBody(
         CompoundStmt::Create(AST, {Assign}, FPOptionsOverride(),
                              SourceLocation(), SourceLocation()));
     Constructor->setAccess(AccessSpecifier::AS_public);
     Record->addDecl(Constructor);
     return *this;
   }
 
   BuiltinTypeDeclBuilder &addArraySubscriptOperators() {
+    if (Record->isCompleteDefinition())
+      return *this;
     addArraySubscriptOperator(true);
     addArraySubscriptOperator(false);
     return *this;
   }
 
   BuiltinTypeDeclBuilder &addArraySubscriptOperator(bool IsConst) {
+    if (Record->isCompleteDefinition())
+      return *this;
     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.");
BuiltinTypeDeclBuilder &addArraySubscriptOperator(bool IsConst) {
     MethodDecl->setLexicalDeclContext(Record);
     MethodDecl->setAccess(AccessSpecifier::AS_public);
     Record->addDecl(MethodDecl);
 
     return *this;
   }
 
   BuiltinTypeDeclBuilder &startDefinition() {
+    if (Record->isCompleteDefinition())
+      return *this;
     Record->startDefinition();
     return *this;
   }
 
   BuiltinTypeDeclBuilder &completeDefinition() {
+    if (Record->isCompleteDefinition())
+      return *this;
     assert(Record->isBeingDefined() &&
            "Definition must be started before completing it.");
 
     Record->completeDefinition();
     return *this;
   }
 
   TemplateParameterListBuilder addTemplateArgumentList();
 };
 
 struct TemplateParameterListBuilder {
   BuiltinTypeDeclBuilder &Builder;
   ASTContext &AST;
   llvm::SmallVector<NamedDecl *> Params;
 
   TemplateParameterListBuilder(BuiltinTypeDeclBuilder &RB)
       : Builder(RB), AST(RB.Record->getASTContext()) {}
 
   ~TemplateParameterListBuilder() { finalizeTemplateArgs(); }
 
   TemplateParameterListBuilder &
   addTypeParameter(StringRef Name, QualType DefaultValue = QualType()) {
+    if (Builder.Record->isCompleteDefinition())
+      return *this;
     unsigned Position = static_cast<unsigned>(Params.size());
     auto *Decl = TemplateTypeParmDecl::Create(
         AST, Builder.Record->getDeclContext(), SourceLocation(),
         SourceLocation(), /* TemplateDepth */ 0, Position,
         &AST.Idents.get(Name, tok::TokenKind::identifier), /* Typename */ false,
         /* ParameterPack */ false);
     if (!DefaultValue.isNull())
       Decl->setDefaultArgument(AST.getTrivialTypeSourceInfo(DefaultValue));
auto *ParamList =
                                       Params, SourceLocation(), nullptr);
     Builder.Template = ClassTemplateDecl::Create(
         AST, Builder.Record->getDeclContext(), SourceLocation(),
         DeclarationName(Builder.Record->getIdentifier()), ParamList,
         Builder.Record);
     Builder.Record->setDescribedClassTemplate(Builder.Template);
     Builder.Template->setImplicit(true);
     Builder.Template->setLexicalDeclContext(Builder.Record->getDeclContext());
+    // NOTE: setPreviousDecl before addDecl so new decl replace old decl when
+    // make visible.
+    Builder.Template->setPreviousDecl(Builder.PrevTemplate);
     Builder.Record->getDeclContext()->addDecl(Builder.Template);
     Params.clear();
 
     QualType T = Builder.Template->getInjectedClassNameSpecialization();
     T = AST.getInjectedClassNameType(Builder.Record, T);
 
     return Builder;
   }
 };
 
 TemplateParameterListBuilder BuiltinTypeDeclBuilder::addTemplateArgumentList() {
   return TemplateParameterListBuilder(*this);
 }
 } // namespace
 
 HLSLExternalSemaSource::~HLSLExternalSemaSource() {}
 
 void HLSLExternalSemaSource::InitializeSema(Sema &S) {
   SemaPtr = &S;
   ASTContext &AST = SemaPtr->getASTContext();
+  // If the translation unit has external storage force external decls to load.
+  if (AST.getTranslationUnitDecl()->hasExternalLexicalStorage())
+    (void)AST.getTranslationUnitDecl()->decls_begin();
+
   IdentifierInfo &HLSL = AST.Idents.get("hlsl", tok::TokenKind::identifier);
-  HLSLNamespace =
-      NamespaceDecl::Create(AST, AST.getTranslationUnitDecl(), false,
-                            SourceLocation(), SourceLocation(), &HLSL, nullptr);
+  LookupResult Result(S, &HLSL, SourceLocation(), Sema::LookupNamespaceName);
+  NamespaceDecl *PrevDecl = nullptr;
+  if (S.LookupQualifiedName(Result, AST.getTranslationUnitDecl()))
+    PrevDecl = Result.getAsSingle<NamespaceDecl>();
+  HLSLNamespace = NamespaceDecl::Create(AST, AST.getTranslationUnitDecl(),
+                                        false, SourceLocation(),
+                                        SourceLocation(), &HLSL, PrevDecl);
   HLSLNamespace->setImplicit(true);
+  HLSLNamespace->setHasExternalLexicalStorage();
   AST.getTranslationUnitDecl()->addDecl(HLSLNamespace);
+
+  // Force external decls in the HLSL namespace to load from the PCH.
+  HLSLNamespace->getCanonicalDecl()->decls_begin();
   defineTrivialHLSLTypes();
   forwardDeclareHLSLTypes();
 
   // This adds a `using namespace hlsl` directive. In DXC, we don't put HLSL's
   // built in types inside a namespace, but we are planning to change that in
   // the near future. In order to be source compatible older versions of HLSL
   // will need to implicitly use the hlsl namespace. For now in clang everything
   // will get added to the namespace, and we can remove the using directive for
   // future language versions to match HLSL's evolution.
   auto *UsingDecl = UsingDirectiveDecl::Create(
       AST, AST.getTranslationUnitDecl(), SourceLocation(), SourceLocation(),
       NestedNameSpecifierLoc(), SourceLocation(), HLSLNamespace,
       AST.getTranslationUnitDecl());
 
   AST.getTranslationUnitDecl()->addDecl(UsingDecl);
 }
 
 void HLSLExternalSemaSource::defineHLSLVectorAlias() {
   ASTContext &AST = SemaPtr->getASTContext();
 
void HLSLExternalSemaSource::defineHLSLVectorAlias() {
   Template->setImplicit(true);
   Template->setLexicalDeclContext(Record->getDeclContext());
   HLSLNamespace->addDecl(Template);
 }
 
 void HLSLExternalSemaSource::defineTrivialHLSLTypes() {
   defineHLSLVectorAlias();
 
+  ASTContext &AST = SemaPtr->getASTContext();
+  IdentifierInfo &II = AST.Idents.get("Resource", tok::TokenKind::identifier);
+  LookupResult Result(*SemaPtr, &II, SourceLocation(), Sema::LookupTagName);
+
   ResourceDecl = BuiltinTypeDeclBuilder(*SemaPtr, HLSLNamespace, "Resource")
                      .startDefinition()
                      .addHandleMember(AccessSpecifier::AS_public)
                      .completeDefinition()
                      .Record;
 }
 
 void HLSLExternalSemaSource::forwardDeclareHLSLTypes() {
   CXXRecordDecl *Decl;
   Decl = BuiltinTypeDeclBuilder(*SemaPtr, HLSLNamespace, "RWBuffer")
              .addTemplateArgumentList()
              .addTypeParameter("element_type", SemaPtr->getASTContext().FloatTy)
              .finalizeTemplateArgs()
              .Record;
-  Completions.insert(std::make_pair(
-      Decl, std::bind(&HLSLExternalSemaSource::completeBufferType, this,
-                      std::placeholders::_1)));
+  if (!Decl->isCompleteDefinition())
+    Completions.insert(
+        std::make_pair(Decl->getCanonicalDecl(),
+                       std::bind(&HLSLExternalSemaSource::completeBufferType,
+                                 this, std::placeholders::_1)));
 }
 
 void HLSLExternalSemaSource::CompleteType(TagDecl *Tag) {
   if (!isa<CXXRecordDecl>(Tag))
     return;
   auto Record = cast<CXXRecordDecl>(Tag);
 
   // If this is a specialization, we need to get the underlying templated
   // declaration and complete that.
   if (auto TDecl = dyn_cast<ClassTemplateSpecializationDecl>(Record))
     Record = TDecl->getSpecializedTemplate()->getTemplatedDecl();
+  Record = Record->getCanonicalDecl();
   auto It = Completions.find(Record);
   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/test/AST/HLSL/Inputs/pch.hlsl

+
+float2 foo(float2 a, float2 b) {
+  return a + b;
+}

clang/test/AST/HLSL/Inputs/pch_with_buf.hlsl

+
+float2 foo(float2 a, float2 b) {
+  return a + b;
+}
+
+RWBuffer<float> Buf;

clang/test/AST/HLSL/pch.hlsl

+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl \
+// RUN:  -finclude-default-header -emit-pch -o %t %S/Inputs/pch.hlsl
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl \
+// RUN:  -finclude-default-header -include-pch %t -fsyntax-only -ast-dump-all %s \
+// RUN: | FileCheck  %s
+
+// Make sure PCH works by using function declared in PCH header and declare a RWBuffer in current file.
+// CHECK:FunctionDecl 0x[[FOO:[0-9a-f]+]] <{{.*}}:2:1, line:4:1> line:2:8 imported used foo 'float2 (float2, float2)'
+// CHECK:VarDecl 0x{{[0-9a-f]+}} <{{.*}}:10:1, col:23> col:23 Buffer 'hlsl::RWBuffer<float>':'hlsl::RWBuffer<>'
+hlsl::RWBuffer<float> Buffer;
+
+float2 bar(float2 a, float2 b) {
+// CHECK:CallExpr 0x{{[0-9a-f]+}} <col:10, col:18> 'float2':'float __attribute__((ext_vector_type(2)))'
+// CHECK-NEXT:ImplicitCastExpr 0x{{[0-9a-f]+}} <col:10> 'float2 (*)(float2, float2)' <FunctionToPointerDecay>
+// CHECK-NEXT:`-DeclRefExpr 0x{{[0-9a-f]+}} <col:10> 'float2 (float2, float2)' lvalue Function 0x[[FOO]] 'foo' 'float2 (float2, float2)'
+  return foo(a, b);
+}

clang/test/AST/HLSL/pch_with_buf.hlsl

+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -finclude-default-header -emit-pch -o %t %S/Inputs/pch_with_buf.hlsl
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl \
+// RUN:  -finclude-default-header -include-pch %t -fsyntax-only -ast-dump-all %s | FileCheck  %s
+
+// Make sure PCH works by using function declared in PCH header.
+// CHECK:FunctionDecl 0x[[FOO:[0-9a-f]+]] <{{.*}}:2:1, line:4:1> line:2:8 imported used foo 'float2 (float2, float2)'
+// Make sure buffer defined in PCH works.
+// CHECK:VarDecl 0x{{[0-9a-f]+}} <line:6:1, col:17> col:17 imported Buf 'RWBuffer<float>':'hlsl::RWBuffer<>'
+// Make sure declare a RWBuffer in current file works.
+// CHECK:VarDecl 0x{{[0-9a-f]+}} <{{.*}}:11:1, col:23> col:23 Buf2 'hlsl::RWBuffer<float>':'hlsl::RWBuffer<>'
+hlsl::RWBuffer<float> Buf2;
+
+float2 bar(float2 a, float2 b) {
+// CHECK:CallExpr 0x{{[0-9a-f]+}} <col:10, col:18> 'float2':'float __attribute__((ext_vector_type(2)))'
+// CHECK-NEXT:ImplicitCastExpr 0x{{[0-9a-f]+}} <col:10> 'float2 (*)(float2, float2)' <FunctionToPointerDecay>
+// CHECK-NEXT:`-DeclRefExpr 0x{{[0-9a-f]+}} <col:10> 'float2 (float2, float2)' lvalue Function 0x[[FOO]] 'foo' 'float2 (float2, float2)'
+  return foo(a, b);
+}