[llvm-branch-commits] [DXIL][Analysis] Implement enough of DXILResourceAnalysis for buffers (PR #100699)

[Xiang Li via llvm-branch-commits]

================
@@ -331,6 +336,249 @@ std::pair<uint32_t, uint32_t> ResourceInfo::getAnnotateProps() const {
   return {Word0, Word1};
 }
 
+void ResourceInfo::print(raw_ostream &OS) const {
+  OS << "  Symbol: ";
+  Symbol->printAsOperand(OS);
+  OS << "\n";
+
+  OS << "  Name: \"" << Name << "\"\n"
+     << "  Binding:\n"
+     << "    Unique ID: " << Binding.UniqueID << "\n"
+     << "    Space: " << Binding.Space << "\n"
+     << "    Lower Bound: " << Binding.LowerBound << "\n"
+     << "    Size: " << Binding.Size << "\n"
+     << "  Class: " << static_cast<unsigned>(RC) << "\n"
+     << "  Kind: " << static_cast<unsigned>(Kind) << "\n";
+
+  if (isCBuffer()) {
+    OS << "  CBuffer size: " << CBufferSize << "\n";
+  } else if (isSampler()) {
+    OS << "  Sampler Type: " << static_cast<unsigned>(SamplerTy) << "\n";
+  } else {
+    if (isUAV()) {
+      OS << "  Globally Coherent: " << UAVFlags.GloballyCoherent << "\n"
+         << "  HasCounter: " << UAVFlags.HasCounter << "\n"
+         << "  IsROV: " << UAVFlags.IsROV << "\n";
+    }
+    if (isMultiSample())
+      OS << "  Sample Count: " << MultiSample.Count << "\n";
+
+    if (isStruct()) {
+      OS << "  Buffer Stride: " << Struct.Stride << "\n";
+      uint32_t AlignLog2 = Struct.Alignment ? Log2(*Struct.Alignment) : 0;
+      OS << "  Alignment: " << AlignLog2 << "\n";
+    } else if (isTyped()) {
+      OS << "  Element Type: " << static_cast<unsigned>(Typed.ElementTy) << "\n"
+         << "  Element Count: " << static_cast<unsigned>(Typed.ElementCount)
+         << "\n";
+    } else if (isFeedback())
+      OS << "  Feedback Type: " << static_cast<unsigned>(Feedback.Type) << "\n";
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// ResourceMapper
+
+static dxil::ElementType toDXILElementType(Type *Ty, bool IsSigned) {
+  // TODO: Handle unorm, snorm, and packed.
+  Ty = Ty->getScalarType();
+
+  if (Ty->isIntegerTy()) {
+    switch (Ty->getIntegerBitWidth()) {
+    case 16:
+      return IsSigned ? ElementType::I16 : ElementType::U16;
+    case 32:
+      return IsSigned ? ElementType::I32 : ElementType::U32;
+    case 64:
+      return IsSigned ? ElementType::I64 : ElementType::U64;
+    case 1:
+    default:
+      return ElementType::Invalid;
+    }
+  } else if (Ty->isFloatTy()) {
+    return ElementType::F32;
+  } else if (Ty->isDoubleTy()) {
+    return ElementType::F64;
+  } else if (Ty->isHalfTy()) {
+    return ElementType::F16;
+  }
+
+  return ElementType::Invalid;
+}
+
+namespace {
+
+class ResourceMapper {
+  Module &M;
+  LLVMContext &Context;
+  DXILResourceMap &Resources;
+
+  // Unique ID is per resource type to match DXC.
+  uint32_t NextUAV = 0;
+  uint32_t NextSRV = 0;
+  uint32_t NextCBuf = 0;
+  uint32_t NextSmp = 0;
+
+public:
+  ResourceMapper(Module &M,
+                 MapVector<CallInst *, dxil::ResourceInfo> &Resources)
+      : M(M), Context(M.getContext()), Resources(Resources) {}
+
+  void diagnoseHandle(CallInst *CI, const Twine &Msg,
+                      DiagnosticSeverity Severity = DS_Error) {
+    std::string S;
+    raw_string_ostream SS(S);
+    CI->printAsOperand(SS);
+    DiagnosticInfoUnsupported Diag(*CI->getFunction(), Msg + ": " + SS.str(),
+                                   CI->getDebugLoc(), Severity);
+    Context.diagnose(Diag);
+  }
+
+  ResourceInfo *mapBufferType(CallInst *CI, TargetExtType *HandleTy,
+                              bool IsTyped) {
+    if (HandleTy->getNumTypeParameters() != 1 ||
+        HandleTy->getNumIntParameters() != (IsTyped ? 3 : 2)) {
+      diagnoseHandle(CI, Twine("Invalid buffer target type"));
+      return nullptr;
+    }
+
+    Type *ElTy = HandleTy->getTypeParameter(0);
+    unsigned IsWriteable = HandleTy->getIntParameter(0);
+    unsigned IsROV = HandleTy->getIntParameter(1);
+    bool IsSigned = IsTyped && HandleTy->getIntParameter(2);
+
+    ResourceClass RC = IsWriteable ? ResourceClass::UAV : ResourceClass::SRV;
+    ResourceKind Kind;
+    if (IsTyped)
+      Kind = ResourceKind::TypedBuffer;
+    else if (ElTy->isIntegerTy(8))
+      Kind = ResourceKind::RawBuffer;
+    else
+      Kind = ResourceKind::StructuredBuffer;
+
+    // TODO: We need to lower to a typed pointer, can we smuggle the type
+    // through?
+    Value *Symbol = UndefValue::get(PointerType::getUnqual(Context));
+    // TODO: We don't actually keep track of the name right now...
+    StringRef Name = "";
+
+    auto [It, Success] = Resources.try_emplace(CI, RC, Kind, Symbol, Name);
+    assert(Success && "Mapping the same CallInst again?");
+    (void)Success;
+    // We grab a pointer into the map's storage, which isn't generally safe.
+    // Since we're just using this to fill in the info the map won't mutate and
+    // the pointer stays valid for as long as we need it to.
+    ResourceInfo *RI = &(It->second);
+
+    if (RI->isUAV())
+      // TODO: We need analysis for GloballyCoherent and HasCounter
+      RI->setUAV(false, false, IsROV);
+
+    if (RI->isTyped()) {
+      dxil::ElementType ET = toDXILElementType(ElTy, IsSigned);
+      uint32_t Count = 1;
+      if (auto *VTy = dyn_cast<FixedVectorType>(ElTy))
+        Count = VTy->getNumElements();
+      RI->setTyped(ET, Count);
+    } else if (RI->isStruct()) {
+      const DataLayout &DL = M.getDataLayout();
+
+      // This mimics what DXC does. Notably, we only ever set the alignment if
+      // the type is actually a struct type.
+      uint32_t Stride = DL.getTypeAllocSize(ElTy);
+      MaybeAlign Alignment;
+      if (auto *STy = dyn_cast<StructType>(ElTy))
+        Alignment = DL.getStructLayout(STy)->getAlignment();
+      RI->setStruct(Stride, Alignment);
+    }
+
+    return RI;
+  }
+
+  ResourceInfo *mapHandleIntrin(CallInst *CI) {
+    FunctionType *FTy = CI->getFunctionType();
+    Type *RetTy = FTy->getReturnType();
+    auto *HandleTy = dyn_cast<TargetExtType>(RetTy);
+    if (!HandleTy) {
+      diagnoseHandle(CI, "dx.handle.fromBinding requires target type");
+      return nullptr;
+    }
+
+    StringRef TypeName = HandleTy->getName();
+    if (TypeName == "dx.TypedBuffer") {
+      return mapBufferType(CI, HandleTy, /*IsTyped=*/true);
+    } else if (TypeName == "dx.RawBuffer") {
+      return mapBufferType(CI, HandleTy, /*IsTyped=*/false);
+    } else if (TypeName == "dx.CBuffer") {
+      // TODO: implement
+      diagnoseHandle(CI, "dx.CBuffer handles are not implemented yet");
----------------
python3kgae wrote:

Could we use ptr with special address space for RawBuffer and CBuffer instead of handles for clang codeGen output?
Then dx.RawBuffer could be dx.StructuredBuffer, since it doesn't share between RawBuffer and StructuredBuffer anymore.

https://github.com/llvm/llvm-project/pull/100699