Skip to content

Commit 02c9dae

Browse files
authored
[HLSL] Add support to lookup a ResourceBindingInfo from its use (#126556)
Adds `findByUse` which takes a `llvm::Value` from a use and resolves it (as best as possible) back to the creation of that resource. It may return multiple ResourceBindingInfo if the use comes from branched control flow. Fixes #125746
1 parent 0cc7381 commit 02c9dae

File tree

4 files changed

+331
-0
lines changed

4 files changed

+331
-0
lines changed

llvm/include/llvm/Analysis/DXILResource.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,13 @@ class DXILBindingMap {
446446
return Pos == CallMap.end() ? Infos.end() : (Infos.begin() + Pos->second);
447447
}
448448

449+
/// Resolves a resource handle into a vector of ResourceBindingInfos that
450+
/// represent the possible unique creations of the handle. Certain cases are
451+
/// ambiguous so multiple creation instructions may be returned. The resulting
452+
/// ResourceBindingInfo can be used to depuplicate unique handles that
453+
/// reference the same resource
454+
SmallVector<dxil::ResourceBindingInfo> findByUse(const Value *Key) const;
455+
449456
const_iterator find(const CallInst *Key) const {
450457
auto Pos = CallMap.find(Key);
451458
return Pos == CallMap.end() ? Infos.end() : (Infos.begin() + Pos->second);

llvm/lib/Analysis/DXILResource.cpp

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,45 @@ void DXILBindingMap::print(raw_ostream &OS, DXILResourceTypeMap &DRTM,
770770
}
771771
}
772772

773+
SmallVector<dxil::ResourceBindingInfo>
774+
DXILBindingMap::findByUse(const Value *Key) const {
775+
if (const PHINode *Phi = dyn_cast<PHINode>(Key)) {
776+
SmallVector<dxil::ResourceBindingInfo> Children;
777+
for (const Value *V : Phi->operands()) {
778+
Children.append(findByUse(V));
779+
}
780+
return Children;
781+
}
782+
783+
const CallInst *CI = dyn_cast<CallInst>(Key);
784+
if (!CI)
785+
return {};
786+
787+
switch (CI->getIntrinsicID()) {
788+
// Found the create, return the binding
789+
case Intrinsic::dx_resource_handlefrombinding: {
790+
const auto *It = find(CI);
791+
assert(It != Infos.end() && "HandleFromBinding must be in resource map");
792+
return {*It};
793+
}
794+
default:
795+
break;
796+
}
797+
798+
// Check if any of the parameters are the resource we are following. If so
799+
// keep searching. If none of them are return an empty list
800+
const Type *UseType = CI->getType();
801+
SmallVector<dxil::ResourceBindingInfo> Children;
802+
for (const Value *V : CI->args()) {
803+
if (V->getType() != UseType)
804+
continue;
805+
806+
Children.append(findByUse(V));
807+
}
808+
809+
return Children;
810+
}
811+
773812
//===----------------------------------------------------------------------===//
774813

775814
AnalysisKey DXILResourceTypeAnalysis::Key;

llvm/unittests/Target/DirectX/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,12 @@ set(LLVM_LINK_COMPONENTS
88
Core
99
DirectXCodeGen
1010
DirectXPointerTypeAnalysis
11+
Passes
1112
Support
1213
)
1314

1415
add_llvm_target_unittest(DirectXTests
1516
CBufferDataLayoutTests.cpp
1617
PointerTypeAnalysisTests.cpp
18+
UniqueResourceFromUseTests.cpp
1719
)
Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
//===- llvm/unittests/Target/DirectX/PointerTypeAnalysisTests.cpp ---------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "DirectXTargetMachine.h"
10+
#include "llvm/Analysis/DXILResource.h"
11+
#include "llvm/AsmParser/Parser.h"
12+
#include "llvm/CodeGen/CommandFlags.h"
13+
#include "llvm/IR/Instructions.h"
14+
#include "llvm/IR/LLVMContext.h"
15+
#include "llvm/IR/Module.h"
16+
#include "llvm/IR/Type.h"
17+
#include "llvm/Passes/PassBuilder.h"
18+
#include "llvm/Support/Casting.h"
19+
#include "llvm/Support/SourceMgr.h"
20+
21+
#include "gtest/gtest.h"
22+
23+
using namespace llvm;
24+
using namespace llvm::dxil;
25+
26+
namespace {
27+
class UniqueResourceFromUseTest : public testing::Test {
28+
protected:
29+
PassBuilder *PB;
30+
ModuleAnalysisManager *MAM;
31+
32+
virtual void SetUp() {
33+
MAM = new ModuleAnalysisManager();
34+
PB = new PassBuilder();
35+
PB->registerModuleAnalyses(*MAM);
36+
MAM->registerPass([&] { return DXILResourceTypeAnalysis(); });
37+
MAM->registerPass([&] { return DXILResourceBindingAnalysis(); });
38+
}
39+
40+
virtual void TearDown() {
41+
delete PB;
42+
delete MAM;
43+
}
44+
};
45+
46+
TEST_F(UniqueResourceFromUseTest, TestTrivialUse) {
47+
StringRef Assembly = R"(
48+
define void @main() {
49+
entry:
50+
%handle = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 1, i32 2, i32 3, i32 4, i1 false)
51+
call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
52+
call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
53+
ret void
54+
}
55+
56+
declare target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32, i32, i32, i32, i1)
57+
declare void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
58+
)";
59+
60+
LLVMContext Context;
61+
SMDiagnostic Error;
62+
auto M = parseAssemblyString(Assembly, Error, Context);
63+
ASSERT_TRUE(M) << "Bad assembly?";
64+
65+
const DXILBindingMap &DBM = MAM->getResult<DXILResourceBindingAnalysis>(*M);
66+
for (const Function &F : M->functions()) {
67+
if (F.getName() != "a.func") {
68+
continue;
69+
}
70+
71+
unsigned CalledResources = 0;
72+
73+
for (const User *U : F.users()) {
74+
const CallInst *CI = cast<CallInst>(U);
75+
const Value *Handle = CI->getArgOperand(0);
76+
const auto Bindings = DBM.findByUse(Handle);
77+
ASSERT_EQ(Bindings.size(), 1u)
78+
<< "Handle should resolve into one resource";
79+
80+
auto Binding = Bindings[0].getBinding();
81+
EXPECT_EQ(0u, Binding.RecordID);
82+
EXPECT_EQ(1u, Binding.Space);
83+
EXPECT_EQ(2u, Binding.LowerBound);
84+
EXPECT_EQ(3u, Binding.Size);
85+
86+
CalledResources++;
87+
}
88+
89+
EXPECT_EQ(2u, CalledResources)
90+
<< "Expected 2 resolved call to create resource";
91+
}
92+
}
93+
94+
TEST_F(UniqueResourceFromUseTest, TestIndirectUse) {
95+
StringRef Assembly = R"(
96+
define void @foo() {
97+
%handle = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 1, i32 2, i32 3, i32 4, i1 false)
98+
%handle2 = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %handle)
99+
%handle3 = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %handle2)
100+
%handle4 = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %handle3)
101+
call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle4)
102+
ret void
103+
}
104+
105+
declare target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32, i32, i32, i32, i1)
106+
declare void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
107+
declare target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %handle)
108+
)";
109+
110+
LLVMContext Context;
111+
SMDiagnostic Error;
112+
auto M = parseAssemblyString(Assembly, Error, Context);
113+
ASSERT_TRUE(M) << "Bad assembly?";
114+
115+
const DXILBindingMap &DBM = MAM->getResult<DXILResourceBindingAnalysis>(*M);
116+
for (const Function &F : M->functions()) {
117+
if (F.getName() != "a.func") {
118+
continue;
119+
}
120+
121+
unsigned CalledResources = 0;
122+
123+
for (const User *U : F.users()) {
124+
const CallInst *CI = cast<CallInst>(U);
125+
const Value *Handle = CI->getArgOperand(0);
126+
const auto Bindings = DBM.findByUse(Handle);
127+
ASSERT_EQ(Bindings.size(), 1u)
128+
<< "Handle should resolve into one resource";
129+
130+
auto Binding = Bindings[0].getBinding();
131+
EXPECT_EQ(0u, Binding.RecordID);
132+
EXPECT_EQ(1u, Binding.Space);
133+
EXPECT_EQ(2u, Binding.LowerBound);
134+
EXPECT_EQ(3u, Binding.Size);
135+
136+
CalledResources++;
137+
}
138+
139+
EXPECT_EQ(1u, CalledResources)
140+
<< "Expected 1 resolved call to create resource";
141+
}
142+
}
143+
144+
TEST_F(UniqueResourceFromUseTest, TestAmbigousIndirectUse) {
145+
StringRef Assembly = R"(
146+
define void @foo() {
147+
%foo = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 1, i32 1, i32 1, i32 1, i1 false)
148+
%bar = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 2, i32 2, i32 2, i32 2, i1 false)
149+
%baz = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 3, i32 3, i32 3, i32 3, i1 false)
150+
%bat = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 4, i32 4, i32 4, i32 4, i1 false)
151+
%a = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %foo, target("dx.RawBuffer", float, 1, 0) %bar)
152+
%b = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %baz, target("dx.RawBuffer", float, 1, 0) %bat)
153+
%handle = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %a, target("dx.RawBuffer", float, 1, 0) %b)
154+
call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
155+
ret void
156+
}
157+
158+
declare target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32, i32, i32, i32, i1)
159+
declare void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
160+
declare target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %x, target("dx.RawBuffer", float, 1, 0) %y)
161+
)";
162+
163+
LLVMContext Context;
164+
SMDiagnostic Error;
165+
auto M = parseAssemblyString(Assembly, Error, Context);
166+
ASSERT_TRUE(M) << "Bad assembly?";
167+
168+
const DXILBindingMap &DBM = MAM->getResult<DXILResourceBindingAnalysis>(*M);
169+
for (const Function &F : M->functions()) {
170+
if (F.getName() != "a.func") {
171+
continue;
172+
}
173+
174+
unsigned CalledResources = 0;
175+
176+
for (const User *U : F.users()) {
177+
const CallInst *CI = cast<CallInst>(U);
178+
const Value *Handle = CI->getArgOperand(0);
179+
const auto Bindings = DBM.findByUse(Handle);
180+
ASSERT_EQ(Bindings.size(), 4u)
181+
<< "Handle should resolve into four resources";
182+
183+
auto Binding = Bindings[0].getBinding();
184+
EXPECT_EQ(0u, Binding.RecordID);
185+
EXPECT_EQ(1u, Binding.Space);
186+
EXPECT_EQ(1u, Binding.LowerBound);
187+
EXPECT_EQ(1u, Binding.Size);
188+
189+
Binding = Bindings[1].getBinding();
190+
EXPECT_EQ(1u, Binding.RecordID);
191+
EXPECT_EQ(2u, Binding.Space);
192+
EXPECT_EQ(2u, Binding.LowerBound);
193+
EXPECT_EQ(2u, Binding.Size);
194+
195+
Binding = Bindings[2].getBinding();
196+
EXPECT_EQ(2u, Binding.RecordID);
197+
EXPECT_EQ(3u, Binding.Space);
198+
EXPECT_EQ(3u, Binding.LowerBound);
199+
EXPECT_EQ(3u, Binding.Size);
200+
201+
Binding = Bindings[3].getBinding();
202+
EXPECT_EQ(3u, Binding.RecordID);
203+
EXPECT_EQ(4u, Binding.Space);
204+
EXPECT_EQ(4u, Binding.LowerBound);
205+
EXPECT_EQ(4u, Binding.Size);
206+
207+
CalledResources++;
208+
}
209+
210+
EXPECT_EQ(1u, CalledResources)
211+
<< "Expected 1 resolved call to create resource";
212+
}
213+
}
214+
215+
TEST_F(UniqueResourceFromUseTest, TestConditionalUse) {
216+
StringRef Assembly = R"(
217+
define void @foo(i32 %n) {
218+
entry:
219+
%x = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 1, i32 1, i32 1, i32 1, i1 false)
220+
%y = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 4, i32 4, i32 4, i32 4, i1 false)
221+
%cond = icmp eq i32 %n, 0
222+
br i1 %cond, label %bb.true, label %bb.false
223+
224+
bb.true:
225+
%handle_t = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %x)
226+
br label %bb.exit
227+
228+
bb.false:
229+
%handle_f = call target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %y)
230+
br label %bb.exit
231+
232+
bb.exit:
233+
%handle = phi target("dx.RawBuffer", float, 1, 0) [ %handle_t, %bb.true ], [ %handle_f, %bb.false ]
234+
call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
235+
ret void
236+
}
237+
238+
declare target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32, i32, i32, i32, i1)
239+
declare void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
240+
declare target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", float, 1, 0) %x)
241+
)";
242+
243+
LLVMContext Context;
244+
SMDiagnostic Error;
245+
auto M = parseAssemblyString(Assembly, Error, Context);
246+
ASSERT_TRUE(M) << "Bad assembly?";
247+
248+
const DXILBindingMap &DBM = MAM->getResult<DXILResourceBindingAnalysis>(*M);
249+
for (const Function &F : M->functions()) {
250+
if (F.getName() != "a.func") {
251+
continue;
252+
}
253+
254+
unsigned CalledResources = 0;
255+
256+
for (const User *U : F.users()) {
257+
const CallInst *CI = cast<CallInst>(U);
258+
const Value *Handle = CI->getArgOperand(0);
259+
const auto Bindings = DBM.findByUse(Handle);
260+
ASSERT_EQ(Bindings.size(), 2u)
261+
<< "Handle should resolve into four resources";
262+
263+
auto Binding = Bindings[0].getBinding();
264+
EXPECT_EQ(0u, Binding.RecordID);
265+
EXPECT_EQ(1u, Binding.Space);
266+
EXPECT_EQ(1u, Binding.LowerBound);
267+
EXPECT_EQ(1u, Binding.Size);
268+
269+
Binding = Bindings[1].getBinding();
270+
EXPECT_EQ(1u, Binding.RecordID);
271+
EXPECT_EQ(4u, Binding.Space);
272+
EXPECT_EQ(4u, Binding.LowerBound);
273+
EXPECT_EQ(4u, Binding.Size);
274+
275+
CalledResources++;
276+
}
277+
278+
EXPECT_EQ(1u, CalledResources)
279+
<< "Expected 1 resolved call to create resource";
280+
}
281+
}
282+
283+
} // namespace

0 commit comments

Comments
 (0)