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

[Justin Bogner 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");
----------------
bogner wrote:

I don't know about SPIR-V, but just using a struct for a cbuffer in LLVM IR opens up a bunch of complexity about the layout. Consider:

```hlsl
cbuffer FunnyLayout { 
   float2 first;
   float arr[2];
   float2 last;
};
```

This is layed out like so in a cbuffer, in 3 rows (excuse the ascii art):
```
 0        4        8       12
 +--------+--------+--------+--------+
 |      first      |        |        |
 +--------+--------+--------+--------+
16       20       24       28
 +--------+--------+--------+--------+
 | arr[0] |        |        |        |
 +--------+--------+--------+--------+
32       36       40       44
 +--------+--------+--------+--------+
 | arr[1] |      last       |        |
 +--------+--------+--------+--------+
```

To access "last", we do a cbuffer load with index 2 and take the second element.

There are two problems with using a pointer to an llvm struct type for this. First, the layout of `{<2 x float>, [2 x float], <2 x float>}` is as follows:
```
0        4        8       12
 +--------+--------+--------+--------+
 |      first      | arr[0] | arr[1] |
 +--------+--------+--------+--------+
16       20
 +--------+--------+
 |      last       |
 +--------+--------+
```

Second, even if you were to construct a type with a bunch of padding so that it's layed out like a cbuffer, we'd have to enforce that all loads and stores are of exactly 16 bytes on a 16 byte boundary. This is trivial if we have an intrinsic for it, but requires extra validation and logic in various places and special knowledge in optimizations if we just use plain loads and stores.

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