Skip to content

Commit ecd2bf7

Browse files
vporpovporpo
authored
[SandboxIR] Add setOperand() and RAUW,RUWIf,RUOW (llvm#98410)
This patch adds the following member functions: - User::setOperand() - User::replaceUsesOfWith() - Value::replaceAllUsesWith() - Value::replaceUsesWithIf()
1 parent 144dae2 commit ecd2bf7

File tree

3 files changed

+176
-2
lines changed

3 files changed

+176
-2
lines changed

‎llvm/include/llvm/SandboxIR/SandboxIR.h

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ class Value {
200200
llvm::Value *Val = nullptr;
201201

202202
friend class Context; // For getting `Val`.
203+
friend class User; // For getting `Val`.
203204

204205
/// All values point to the context.
205206
Context &Ctx;
@@ -284,6 +285,11 @@ class Value {
284285
Type *getType() const { return Val->getType(); }
285286

286287
Context &getContext() const { return Ctx; }
288+
289+
void replaceUsesWithIf(Value *OtherV,
290+
llvm::function_ref<bool(const Use &)> ShouldReplace);
291+
void replaceAllUsesWith(Value *Other);
292+
287293
#ifndef NDEBUG
288294
/// Should crash if there is something wrong with the instruction.
289295
virtual void verify() const = 0;
@@ -349,6 +355,10 @@ class User : public Value {
349355
virtual unsigned getUseOperandNo(const Use &Use) const = 0;
350356
friend unsigned Use::getOperandNo() const; // For getUseOperandNo()
351357

358+
#ifndef NDEBUG
359+
void verifyUserOfLLVMUse(const llvm::Use &Use) const;
360+
#endif // NDEBUG
361+
352362
public:
353363
/// For isa/dyn_cast.
354364
static bool classof(const Value *From);
@@ -387,6 +397,11 @@ class User : public Value {
387397
return isa<llvm::User>(Val) ? cast<llvm::User>(Val)->getNumOperands() : 0;
388398
}
389399

400+
virtual void setOperand(unsigned OperandIdx, Value *Operand);
401+
/// Replaces any operands that match \p FromV with \p ToV. Returns whether any
402+
/// operands were replaced.
403+
bool replaceUsesOfWith(Value *FromV, Value *ToV);
404+
390405
#ifndef NDEBUG
391406
void verify() const override {
392407
assert(isa<llvm::User>(Val) && "Expected User!");

‎llvm/lib/SandboxIR/SandboxIR.cpp

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,25 @@ Value::user_iterator Value::user_begin() {
103103

104104
unsigned Value::getNumUses() const { return range_size(Val->users()); }
105105

106+
void Value::replaceUsesWithIf(
107+
Value *OtherV, llvm::function_ref<bool(const Use &)> ShouldReplace) {
108+
assert(getType() == OtherV->getType() && "Can't replace with different type");
109+
llvm::Value *OtherVal = OtherV->Val;
110+
Val->replaceUsesWithIf(
111+
OtherVal, [&ShouldReplace, this](llvm::Use &LLVMUse) -> bool {
112+
User *DstU = cast_or_null<User>(Ctx.getValue(LLVMUse.getUser()));
113+
if (DstU == nullptr)
114+
return false;
115+
return ShouldReplace(Use(&LLVMUse, DstU, Ctx));
116+
});
117+
}
118+
119+
void Value::replaceAllUsesWith(Value *Other) {
120+
assert(getType() == Other->getType() &&
121+
"Replacing with Value of different type!");
122+
Val->replaceAllUsesWith(Other->Val);
123+
}
124+
106125
#ifndef NDEBUG
107126
std::string Value::getName() const {
108127
std::stringstream SS;
@@ -165,6 +184,13 @@ Use User::getOperandUseDefault(unsigned OpIdx, bool Verify) const {
165184
return Use(LLVMUse, const_cast<User *>(this), Ctx);
166185
}
167186

187+
#ifndef NDEBUG
188+
void User::verifyUserOfLLVMUse(const llvm::Use &Use) const {
189+
assert(Ctx.getValue(Use.getUser()) == this &&
190+
"Use not found in this SBUser's operands!");
191+
}
192+
#endif
193+
168194
bool User::classof(const Value *From) {
169195
switch (From->getSubclassID()) {
170196
#define DEF_VALUE(ID, CLASS)
@@ -180,6 +206,15 @@ bool User::classof(const Value *From) {
180206
}
181207
}
182208

209+
void User::setOperand(unsigned OperandIdx, Value *Operand) {
210+
assert(isa<llvm::User>(Val) && "No operands!");
211+
cast<llvm::User>(Val)->setOperand(OperandIdx, Operand->Val);
212+
}
213+
214+
bool User::replaceUsesOfWith(Value *FromV, Value *ToV) {
215+
return cast<llvm::User>(Val)->replaceUsesOfWith(FromV->Val, ToV->Val);
216+
}
217+
183218
#ifndef NDEBUG
184219
void User::dumpCommonHeader(raw_ostream &OS) const {
185220
Value::dumpCommonHeader(OS);
@@ -325,10 +360,11 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
325360
return It->second.get();
326361

327362
if (auto *C = dyn_cast<llvm::Constant>(LLVMV)) {
363+
It->second = std::unique_ptr<Constant>(new Constant(C, *this));
364+
auto *NewC = It->second.get();
328365
for (llvm::Value *COp : C->operands())
329366
getOrCreateValueInternal(COp, C);
330-
It->second = std::unique_ptr<Constant>(new Constant(C, *this));
331-
return It->second.get();
367+
return NewC;
332368
}
333369
if (auto *Arg = dyn_cast<llvm::Argument>(LLVMV)) {
334370
It->second = std::unique_ptr<Argument>(new Argument(Arg, *this));

‎llvm/unittests/SandboxIR/SandboxIRTest.cpp

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,9 @@ define i32 @foo(i32 %v0, i32 %v1) {
122122
BasicBlock *LLVMBB = &*LLVMF.begin();
123123
auto LLVMBBIt = LLVMBB->begin();
124124
Instruction *LLVMI0 = &*LLVMBBIt++;
125+
Instruction *LLVMRet = &*LLVMBBIt++;
126+
Argument *LLVMArg0 = LLVMF.getArg(0);
127+
Argument *LLVMArg1 = LLVMF.getArg(1);
125128

126129
auto &F = *Ctx.createFunction(&LLVMF);
127130
auto &BB = *F.begin();
@@ -203,6 +206,126 @@ OperandNo: 0
203206
EXPECT_FALSE(I0->hasNUses(0u));
204207
EXPECT_TRUE(I0->hasNUses(1u));
205208
EXPECT_FALSE(I0->hasNUses(2u));
209+
210+
// Check User.setOperand().
211+
Ret->setOperand(0, Arg0);
212+
EXPECT_EQ(Ret->getOperand(0), Arg0);
213+
EXPECT_EQ(Ret->getOperandUse(0).get(), Arg0);
214+
EXPECT_EQ(LLVMRet->getOperand(0), LLVMArg0);
215+
216+
Ret->setOperand(0, Arg1);
217+
EXPECT_EQ(Ret->getOperand(0), Arg1);
218+
EXPECT_EQ(Ret->getOperandUse(0).get(), Arg1);
219+
EXPECT_EQ(LLVMRet->getOperand(0), LLVMArg1);
220+
}
221+
222+
TEST_F(SandboxIRTest, RUOW) {
223+
parseIR(C, R"IR(
224+
declare void @bar0()
225+
declare void @bar1()
226+
227+
@glob0 = global ptr @bar0
228+
@glob1 = global ptr @bar1
229+
230+
define i32 @foo(i32 %arg0, i32 %arg1) {
231+
%add0 = add i32 %arg0, %arg1
232+
%gep1 = getelementptr i8, ptr @glob0, i32 1
233+
%gep2 = getelementptr i8, ptr @glob1, i32 1
234+
ret i32 %add0
235+
}
236+
)IR");
237+
llvm::Function &LLVMF = *M->getFunction("foo");
238+
sandboxir::Context Ctx(C);
239+
240+
auto &F = *Ctx.createFunction(&LLVMF);
241+
auto &BB = *F.begin();
242+
auto *Arg0 = F.getArg(0);
243+
auto *Arg1 = F.getArg(1);
244+
auto It = BB.begin();
245+
auto *I0 = &*It++;
246+
auto *I1 = &*It++;
247+
auto *I2 = &*It++;
248+
auto *Ret = &*It++;
249+
250+
bool Replaced;
251+
// Try to replace an operand that doesn't match.
252+
Replaced = I0->replaceUsesOfWith(Ret, Arg1);
253+
EXPECT_FALSE(Replaced);
254+
EXPECT_EQ(I0->getOperand(0), Arg0);
255+
EXPECT_EQ(I0->getOperand(1), Arg1);
256+
257+
// Replace I0 operands when operands differ.
258+
Replaced = I0->replaceUsesOfWith(Arg0, Arg1);
259+
EXPECT_TRUE(Replaced);
260+
EXPECT_EQ(I0->getOperand(0), Arg1);
261+
EXPECT_EQ(I0->getOperand(1), Arg1);
262+
263+
// Replace I0 operands when operands are the same.
264+
Replaced = I0->replaceUsesOfWith(Arg1, Arg0);
265+
EXPECT_TRUE(Replaced);
266+
EXPECT_EQ(I0->getOperand(0), Arg0);
267+
EXPECT_EQ(I0->getOperand(1), Arg0);
268+
269+
// Replace Ret operand.
270+
Replaced = Ret->replaceUsesOfWith(I0, Arg0);
271+
EXPECT_TRUE(Replaced);
272+
EXPECT_EQ(Ret->getOperand(0), Arg0);
273+
274+
// Check RAUW on constant.
275+
auto *Glob0 = cast<sandboxir::Constant>(I1->getOperand(0));
276+
auto *Glob1 = cast<sandboxir::Constant>(I2->getOperand(0));
277+
auto *Glob0Op = Glob0->getOperand(0);
278+
Glob0->replaceUsesOfWith(Glob0Op, Glob1);
279+
EXPECT_EQ(Glob0->getOperand(0), Glob1);
280+
}
281+
282+
TEST_F(SandboxIRTest, RAUW_RUWIf) {
283+
parseIR(C, R"IR(
284+
define void @foo(ptr %ptr) {
285+
%ld0 = load float, ptr %ptr
286+
%ld1 = load float, ptr %ptr
287+
store float %ld0, ptr %ptr
288+
store float %ld0, ptr %ptr
289+
ret void
290+
}
291+
)IR");
292+
llvm::Function &LLVMF = *M->getFunction("foo");
293+
sandboxir::Context Ctx(C);
294+
llvm::BasicBlock *LLVMBB = &*LLVMF.begin();
295+
296+
Ctx.createFunction(&LLVMF);
297+
auto *BB = cast<sandboxir::BasicBlock>(Ctx.getValue(LLVMBB));
298+
auto It = BB->begin();
299+
sandboxir::Instruction *Ld0 = &*It++;
300+
sandboxir::Instruction *Ld1 = &*It++;
301+
sandboxir::Instruction *St0 = &*It++;
302+
sandboxir::Instruction *St1 = &*It++;
303+
// Check RUWIf when the lambda returns false.
304+
Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return false; });
305+
EXPECT_EQ(St0->getOperand(0), Ld0);
306+
EXPECT_EQ(St1->getOperand(0), Ld0);
307+
// Check RUWIf when the lambda returns true.
308+
Ld0->replaceUsesWithIf(Ld1, [](const sandboxir::Use &Use) { return true; });
309+
EXPECT_EQ(St0->getOperand(0), Ld1);
310+
EXPECT_EQ(St1->getOperand(0), Ld1);
311+
St0->setOperand(0, Ld0);
312+
St1->setOperand(0, Ld0);
313+
// Check RUWIf user == St0.
314+
Ld0->replaceUsesWithIf(
315+
Ld1, [St0](const sandboxir::Use &Use) { return Use.getUser() == St0; });
316+
EXPECT_EQ(St0->getOperand(0), Ld1);
317+
EXPECT_EQ(St1->getOperand(0), Ld0);
318+
St0->setOperand(0, Ld0);
319+
// Check RUWIf user == St1.
320+
Ld0->replaceUsesWithIf(
321+
Ld1, [St1](const sandboxir::Use &Use) { return Use.getUser() == St1; });
322+
EXPECT_EQ(St0->getOperand(0), Ld0);
323+
EXPECT_EQ(St1->getOperand(0), Ld1);
324+
St1->setOperand(0, Ld0);
325+
// Check RAUW.
326+
Ld1->replaceAllUsesWith(Ld0);
327+
EXPECT_EQ(St0->getOperand(0), Ld0);
328+
EXPECT_EQ(St1->getOperand(0), Ld0);
206329
}
207330

208331
// Check that the operands/users are counted correctly.

0 commit comments

Comments
 (0)