Skip to content

Commit 618b0b7

Browse files
authored
[SandboxIR] Add more Instruction member functions (#98588)
This patch adds new Instruction member functions, including: - getNextNode() - getPrevNode() - getOpcode() - removeFromParent() - eraseFromParent() - getParent() - insertBefore() - insertAfter() - insertInto() - moveBefore() - moveAfter()
1 parent 99153c8 commit 618b0b7

File tree

3 files changed

+271
-2
lines changed

3 files changed

+271
-2
lines changed

llvm/include/llvm/SandboxIR/SandboxIR.h

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,9 @@ namespace llvm {
6868

6969
namespace sandboxir {
7070

71-
class Function;
71+
class BasicBlock;
7272
class Context;
73+
class Function;
7374
class Instruction;
7475
class User;
7576
class Value;
@@ -508,6 +509,14 @@ class Instruction : public sandboxir::User {
508509

509510
Opcode Opc;
510511

512+
/// A SandboxIR Instruction may map to multiple LLVM IR Instruction. This
513+
/// returns its topmost LLVM IR instruction.
514+
llvm::Instruction *getTopmostLLVMInstruction() const;
515+
516+
/// \Returns the LLVM IR Instructions that this SandboxIR maps to in program
517+
/// order.
518+
virtual SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const = 0;
519+
511520
public:
512521
static const char *getOpcodeName(Opcode Opc);
513522
#ifndef NDEBUG
@@ -518,6 +527,40 @@ class Instruction : public sandboxir::User {
518527
#endif
519528
/// This is used by BasicBlock::iterator.
520529
virtual unsigned getNumOfIRInstrs() const = 0;
530+
/// \Returns a BasicBlock::iterator for this Instruction.
531+
BBIterator getIterator() const;
532+
/// \Returns the next sandboxir::Instruction in the block, or nullptr if at
533+
/// the end of the block.
534+
Instruction *getNextNode() const;
535+
/// \Returns the previous sandboxir::Instruction in the block, or nullptr if
536+
/// at the beginning of the block.
537+
Instruction *getPrevNode() const;
538+
/// \Returns this Instruction's opcode. Note that SandboxIR has its own opcode
539+
/// state to allow for new SandboxIR-specific instructions.
540+
Opcode getOpcode() const { return Opc; }
541+
/// Detach this from its parent BasicBlock without deleting it.
542+
void removeFromParent();
543+
/// Detach this Value from its parent and delete it.
544+
void eraseFromParent();
545+
/// Insert this detached instruction before \p BeforeI.
546+
void insertBefore(Instruction *BeforeI);
547+
/// Insert this detached instruction after \p AfterI.
548+
void insertAfter(Instruction *AfterI);
549+
/// Insert this detached instruction into \p BB at \p WhereIt.
550+
void insertInto(BasicBlock *BB, const BBIterator &WhereIt);
551+
/// Move this instruction to \p WhereIt.
552+
void moveBefore(BasicBlock &BB, const BBIterator &WhereIt);
553+
/// Move this instruction before \p Before.
554+
void moveBefore(Instruction *Before) {
555+
moveBefore(*Before->getParent(), Before->getIterator());
556+
}
557+
/// Move this instruction after \p After.
558+
void moveAfter(Instruction *After) {
559+
moveBefore(*After->getParent(), std::next(After->getIterator()));
560+
}
561+
/// \Returns the BasicBlock containing this Instruction, or null if it is
562+
/// detached.
563+
BasicBlock *getParent() const;
521564
/// For isa/dyn_cast.
522565
static bool classof(const sandboxir::Value *From);
523566

@@ -543,6 +586,9 @@ class OpaqueInst : public sandboxir::Instruction {
543586
Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
544587
return getOperandUseDefault(OpIdx, Verify);
545588
}
589+
SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const final {
590+
return {cast<llvm::Instruction>(Val)};
591+
}
546592

547593
public:
548594
static bool classof(const sandboxir::Value *From) {
@@ -570,7 +616,8 @@ class BasicBlock : public Value {
570616
/// Builds a graph that contains all values in \p BB in their original form
571617
/// i.e., no vectorization is taking place here.
572618
void buildBasicBlockFromLLVMIR(llvm::BasicBlock *LLVMBB);
573-
friend class Context; // For `buildBasicBlockFromIR`
619+
friend class Context; // For `buildBasicBlockFromIR`
620+
friend class Instruction; // For LLVM Val.
574621

575622
BasicBlock(llvm::BasicBlock *BB, Context &SBCtx)
576623
: Value(ClassID::Block, BB, SBCtx) {
@@ -623,6 +670,12 @@ class Context {
623670
DenseMap<llvm::Value *, std::unique_ptr<sandboxir::Value>>
624671
LLVMValueToValueMap;
625672

673+
/// Remove \p V from the maps and returns the unique_ptr.
674+
std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);
675+
/// Remove \p SBV from all SandboxIR maps and stop owning it. This effectively
676+
/// detaches \p V from the underlying IR.
677+
std::unique_ptr<Value> detach(Value *V);
678+
friend void Instruction::eraseFromParent(); // For detach().
626679
/// Take ownership of VPtr and store it in `LLVMValueToValueMap`.
627680
Value *registerValue(std::unique_ptr<Value> &&VPtr);
628681

llvm/lib/SandboxIR/SandboxIR.cpp

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,115 @@ const char *Instruction::getOpcodeName(Opcode Opc) {
262262
llvm_unreachable("Unknown Opcode");
263263
}
264264

265+
llvm::Instruction *Instruction::getTopmostLLVMInstruction() const {
266+
Instruction *Prev = getPrevNode();
267+
if (Prev == nullptr) {
268+
// If at top of the BB, return the first BB instruction.
269+
return &*cast<llvm::BasicBlock>(getParent()->Val)->begin();
270+
}
271+
// Else get the Previous sandbox IR instruction's bottom IR instruction and
272+
// return its successor.
273+
llvm::Instruction *PrevBotI = cast<llvm::Instruction>(Prev->Val);
274+
return PrevBotI->getNextNode();
275+
}
276+
277+
BBIterator Instruction::getIterator() const {
278+
auto *I = cast<llvm::Instruction>(Val);
279+
return BasicBlock::iterator(I->getParent(), I->getIterator(), &Ctx);
280+
}
281+
282+
Instruction *Instruction::getNextNode() const {
283+
assert(getParent() != nullptr && "Detached!");
284+
assert(getIterator() != getParent()->end() && "Already at end!");
285+
auto *LLVMI = cast<llvm::Instruction>(Val);
286+
assert(LLVMI->getParent() != nullptr && "LLVM IR instr is detached!");
287+
auto *NextLLVMI = LLVMI->getNextNode();
288+
auto *NextI = cast_or_null<Instruction>(Ctx.getValue(NextLLVMI));
289+
if (NextI == nullptr)
290+
return nullptr;
291+
return NextI;
292+
}
293+
294+
Instruction *Instruction::getPrevNode() const {
295+
assert(getParent() != nullptr && "Detached!");
296+
auto It = getIterator();
297+
if (It != getParent()->begin())
298+
return std::prev(getIterator()).get();
299+
return nullptr;
300+
}
301+
302+
void Instruction::removeFromParent() {
303+
// Detach all the LLVM IR instructions from their parent BB.
304+
for (llvm::Instruction *I : getLLVMInstrs())
305+
I->removeFromParent();
306+
}
307+
308+
void Instruction::eraseFromParent() {
309+
assert(users().empty() && "Still connected to users, can't erase!");
310+
// We don't have Tracking yet, so just erase the LLVM IR instructions.
311+
// Erase in reverse to avoid erasing nstructions with attached uses.
312+
for (llvm::Instruction *I : reverse(getLLVMInstrs()))
313+
I->eraseFromParent();
314+
}
315+
316+
void Instruction::moveBefore(BasicBlock &BB, const BBIterator &WhereIt) {
317+
if (std::next(getIterator()) == WhereIt)
318+
// Destination is same as origin, nothing to do.
319+
return;
320+
auto *LLVMBB = cast<llvm::BasicBlock>(BB.Val);
321+
llvm::BasicBlock::iterator It;
322+
if (WhereIt == BB.end()) {
323+
It = LLVMBB->end();
324+
} else {
325+
Instruction *WhereI = &*WhereIt;
326+
It = WhereI->getTopmostLLVMInstruction()->getIterator();
327+
}
328+
// TODO: Move this to the verifier of sandboxir::Instruction.
329+
assert(is_sorted(getLLVMInstrs(),
330+
[](auto *I1, auto *I2) { return I1->comesBefore(I2); }) &&
331+
"Expected program order!");
332+
// Do the actual move in LLVM IR.
333+
for (auto *I : getLLVMInstrs())
334+
I->moveBefore(*LLVMBB, It);
335+
}
336+
337+
void Instruction::insertBefore(Instruction *BeforeI) {
338+
llvm::Instruction *BeforeTopI = BeforeI->getTopmostLLVMInstruction();
339+
// TODO: Move this to the verifier of sandboxir::Instruction.
340+
assert(is_sorted(getLLVMInstrs(),
341+
[](auto *I1, auto *I2) { return I1->comesBefore(I2); }) &&
342+
"Expected program order!");
343+
for (llvm::Instruction *I : getLLVMInstrs())
344+
I->insertBefore(BeforeTopI);
345+
}
346+
347+
void Instruction::insertAfter(Instruction *AfterI) {
348+
insertInto(AfterI->getParent(), std::next(AfterI->getIterator()));
349+
}
350+
351+
void Instruction::insertInto(BasicBlock *BB, const BBIterator &WhereIt) {
352+
llvm::BasicBlock *LLVMBB = cast<llvm::BasicBlock>(BB->Val);
353+
llvm::Instruction *LLVMBeforeI;
354+
llvm::BasicBlock::iterator LLVMBeforeIt;
355+
if (WhereIt != BB->end()) {
356+
Instruction *BeforeI = &*WhereIt;
357+
LLVMBeforeI = BeforeI->getTopmostLLVMInstruction();
358+
LLVMBeforeIt = LLVMBeforeI->getIterator();
359+
} else {
360+
LLVMBeforeI = nullptr;
361+
LLVMBeforeIt = LLVMBB->end();
362+
}
363+
for (llvm::Instruction *I : getLLVMInstrs())
364+
I->insertInto(LLVMBB, LLVMBeforeIt);
365+
}
366+
367+
BasicBlock *Instruction::getParent() const {
368+
auto *BB = cast<llvm::Instruction>(Val)->getParent();
369+
if (BB == nullptr)
370+
return nullptr;
371+
return cast<BasicBlock>(Ctx.getValue(BB));
372+
}
373+
265374
bool Instruction::classof(const sandboxir::Value *From) {
266375
switch (From->getSubclassID()) {
267376
#define DEF_INSTR(ID, OPC, CLASS) \
@@ -344,6 +453,24 @@ BasicBlock::iterator::getInstr(llvm::BasicBlock::iterator It) const {
344453
return cast_or_null<Instruction>(Ctx->getValue(&*It));
345454
}
346455

456+
std::unique_ptr<Value> Context::detachLLVMValue(llvm::Value *V) {
457+
std::unique_ptr<Value> Erased;
458+
auto It = LLVMValueToValueMap.find(V);
459+
if (It != LLVMValueToValueMap.end()) {
460+
auto *Val = It->second.release();
461+
Erased = std::unique_ptr<Value>(Val);
462+
LLVMValueToValueMap.erase(It);
463+
}
464+
return Erased;
465+
}
466+
467+
std::unique_ptr<Value> Context::detach(Value *V) {
468+
assert(V->getSubclassID() != Value::ClassID::Constant &&
469+
"Can't detach a constant!");
470+
assert(V->getSubclassID() != Value::ClassID::User && "Can't detach a user!");
471+
return detachLLVMValue(V->Val);
472+
}
473+
347474
Value *Context::registerValue(std::unique_ptr<Value> &&VPtr) {
348475
assert(VPtr->getSubclassID() != Value::ClassID::User &&
349476
"Can't register a user!");

llvm/unittests/SandboxIR/SandboxIRTest.cpp

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -471,3 +471,92 @@ define void @foo(i32 %v1) {
471471
}
472472
#endif // NDEBUG
473473
}
474+
475+
TEST_F(SandboxIRTest, Instruction) {
476+
parseIR(C, R"IR(
477+
define void @foo(i8 %v1) {
478+
%add0 = add i8 %v1, %v1
479+
%sub1 = sub i8 %add0, %v1
480+
ret void
481+
}
482+
)IR");
483+
llvm::Function *LLVMF = &*M->getFunction("foo");
484+
sandboxir::Context Ctx(C);
485+
sandboxir::Function *F = Ctx.createFunction(LLVMF);
486+
auto *Arg = F->getArg(0);
487+
auto *BB = &*F->begin();
488+
auto It = BB->begin();
489+
auto *I0 = &*It++;
490+
auto *I1 = &*It++;
491+
auto *Ret = &*It++;
492+
493+
// Check getPrevNode().
494+
EXPECT_EQ(Ret->getPrevNode(), I1);
495+
EXPECT_EQ(I1->getPrevNode(), I0);
496+
EXPECT_EQ(I0->getPrevNode(), nullptr);
497+
498+
// Check getNextNode().
499+
EXPECT_EQ(I0->getNextNode(), I1);
500+
EXPECT_EQ(I1->getNextNode(), Ret);
501+
EXPECT_EQ(Ret->getNextNode(), nullptr);
502+
503+
// Check getIterator().
504+
EXPECT_EQ(I0->getIterator(), std::next(BB->begin(), 0));
505+
EXPECT_EQ(I1->getIterator(), std::next(BB->begin(), 1));
506+
EXPECT_EQ(Ret->getIterator(), std::next(BB->begin(), 2));
507+
508+
// Check getOpcode().
509+
EXPECT_EQ(I0->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
510+
EXPECT_EQ(I1->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
511+
EXPECT_EQ(Ret->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
512+
513+
// Check moveBefore(I).
514+
I1->moveBefore(I0);
515+
EXPECT_EQ(I0->getPrevNode(), I1);
516+
EXPECT_EQ(I1->getNextNode(), I0);
517+
518+
// Check moveAfter(I).
519+
I1->moveAfter(I0);
520+
EXPECT_EQ(I0->getNextNode(), I1);
521+
EXPECT_EQ(I1->getPrevNode(), I0);
522+
523+
// Check moveBefore(BB, It).
524+
I1->moveBefore(*BB, BB->begin());
525+
EXPECT_EQ(I1->getPrevNode(), nullptr);
526+
EXPECT_EQ(I1->getNextNode(), I0);
527+
I1->moveBefore(*BB, BB->end());
528+
EXPECT_EQ(I1->getNextNode(), nullptr);
529+
EXPECT_EQ(Ret->getNextNode(), I1);
530+
I1->moveBefore(*BB, std::next(BB->begin()));
531+
EXPECT_EQ(I0->getNextNode(), I1);
532+
EXPECT_EQ(I1->getNextNode(), Ret);
533+
534+
// Check removeFromParent().
535+
I0->removeFromParent();
536+
#ifndef NDEBUG
537+
EXPECT_DEATH(I0->getPrevNode(), ".*Detached.*");
538+
EXPECT_DEATH(I0->getNextNode(), ".*Detached.*");
539+
#endif // NDEBUG
540+
EXPECT_EQ(I0->getParent(), nullptr);
541+
EXPECT_EQ(I1->getPrevNode(), nullptr);
542+
EXPECT_EQ(I0->getOperand(0), Arg);
543+
544+
// Check insertBefore().
545+
I0->insertBefore(I1);
546+
EXPECT_EQ(I1->getPrevNode(), I0);
547+
548+
// Check insertInto().
549+
I0->removeFromParent();
550+
I0->insertInto(BB, BB->end());
551+
EXPECT_EQ(Ret->getNextNode(), I0);
552+
I0->moveBefore(I1);
553+
EXPECT_EQ(I0->getNextNode(), I1);
554+
555+
// Check eraseFromParent().
556+
#ifndef NDEBUG
557+
EXPECT_DEATH(I0->eraseFromParent(), "Still connected to users.*");
558+
#endif
559+
I1->eraseFromParent();
560+
EXPECT_EQ(I0->getNumUses(), 0u);
561+
EXPECT_EQ(I0->getNextNode(), Ret);
562+
}

0 commit comments

Comments
 (0)