Skip to content

Commit 7467f41

Browse files
vporpovporpo
authored
[SandboxIR] Implement ReturnInst (#99784)
This patch adds the implementation of the SandboxIR ReturnInst which mirrors llvm::ReturnInst.
1 parent 173514f commit 7467f41

File tree

4 files changed

+150
-16
lines changed

4 files changed

+150
-16
lines changed

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

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
// |
4141
// +- PHINode
4242
// |
43-
// +- RetInst
43+
// +- ReturnInst
4444
// |
4545
// +- SelectInst
4646
// |
@@ -76,6 +76,7 @@ class Context;
7676
class Function;
7777
class Instruction;
7878
class LoadInst;
79+
class ReturnInst;
7980
class StoreInst;
8081
class User;
8182
class Value;
@@ -173,11 +174,12 @@ class Value {
173174
/// order.
174175
llvm::Value *Val = nullptr;
175176

176-
friend class Context; // For getting `Val`.
177-
friend class User; // For getting `Val`.
178-
friend class Use; // For getting `Val`.
179-
friend class LoadInst; // For getting `Val`.
180-
friend class StoreInst; // For getting `Val`.
177+
friend class Context; // For getting `Val`.
178+
friend class User; // For getting `Val`.
179+
friend class Use; // For getting `Val`.
180+
friend class LoadInst; // For getting `Val`.
181+
friend class StoreInst; // For getting `Val`.
182+
friend class ReturnInst; // For getting `Val`.
181183

182184
/// All values point to the context.
183185
Context &Ctx;
@@ -497,8 +499,9 @@ class Instruction : public sandboxir::User {
497499
/// A SandboxIR Instruction may map to multiple LLVM IR Instruction. This
498500
/// returns its topmost LLVM IR instruction.
499501
llvm::Instruction *getTopmostLLVMInstruction() const;
500-
friend class LoadInst; // For getTopmostLLVMInstruction().
501-
friend class StoreInst; // For getTopmostLLVMInstruction().
502+
friend class LoadInst; // For getTopmostLLVMInstruction().
503+
friend class StoreInst; // For getTopmostLLVMInstruction().
504+
friend class ReturnInst; // For getTopmostLLVMInstruction().
502505

503506
/// \Returns the LLVM IR Instructions that this SandboxIR maps to in program
504507
/// order.
@@ -639,6 +642,43 @@ class StoreInst final : public Instruction {
639642
#endif
640643
};
641644

645+
class ReturnInst final : public Instruction {
646+
/// Use ReturnInst::create() instead of calling the constructor.
647+
ReturnInst(llvm::Instruction *I, Context &Ctx)
648+
: Instruction(ClassID::Ret, Opcode::Ret, I, Ctx) {}
649+
ReturnInst(ClassID SubclassID, llvm::Instruction *I, Context &Ctx)
650+
: Instruction(SubclassID, Opcode::Ret, I, Ctx) {}
651+
friend class Context; // For accessing the constructor in create*()
652+
Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
653+
return getOperandUseDefault(OpIdx, Verify);
654+
}
655+
SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const final {
656+
return {cast<llvm::Instruction>(Val)};
657+
}
658+
static ReturnInst *createCommon(Value *RetVal, IRBuilder<> &Builder,
659+
Context &Ctx);
660+
661+
public:
662+
static ReturnInst *create(Value *RetVal, Instruction *InsertBefore,
663+
Context &Ctx);
664+
static ReturnInst *create(Value *RetVal, BasicBlock *InsertAtEnd,
665+
Context &Ctx);
666+
static bool classof(const Value *From) {
667+
return From->getSubclassID() == ClassID::Ret;
668+
}
669+
unsigned getUseOperandNo(const Use &Use) const final {
670+
return getUseOperandNoDefault(Use);
671+
}
672+
unsigned getNumOfIRInstrs() const final { return 1u; }
673+
/// \Returns null if there is no return value.
674+
Value *getReturnValue() const;
675+
#ifndef NDEBUG
676+
void verify() const final {}
677+
void dump(raw_ostream &OS) const override;
678+
LLVM_DUMP_METHOD void dump() const override;
679+
#endif
680+
};
681+
642682
/// An LLLVM Instruction that has no SandboxIR equivalent class gets mapped to
643683
/// an OpaqueInstr.
644684
class OpaqueInst : public sandboxir::Instruction {
@@ -776,6 +816,8 @@ class Context {
776816
friend LoadInst; // For createLoadInst()
777817
StoreInst *createStoreInst(llvm::StoreInst *SI);
778818
friend StoreInst; // For createStoreInst()
819+
ReturnInst *createReturnInst(llvm::ReturnInst *I);
820+
friend ReturnInst; // For createReturnInst()
779821

780822
public:
781823
Context(LLVMContext &LLVMCtx)

‎llvm/include/llvm/SandboxIR/SandboxIRValues.def

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ DEF_USER(Constant, Constant)
2727
DEF_INSTR(Opaque, OP(Opaque), OpaqueInst)
2828
DEF_INSTR(Load, OP(Load), LoadInst)
2929
DEF_INSTR(Store, OP(Store), StoreInst)
30+
DEF_INSTR(Ret, OP(Ret), ReturnInst)
3031

3132
#ifdef DEF_VALUE
3233
#undef DEF_VALUE

‎llvm/lib/SandboxIR/SandboxIR.cpp

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,48 @@ void StoreInst::dump() const {
540540
dump(dbgs());
541541
dbgs() << "\n";
542542
}
543+
#endif // NDEBUG
544+
545+
ReturnInst *ReturnInst::createCommon(Value *RetVal, IRBuilder<> &Builder,
546+
Context &Ctx) {
547+
llvm::ReturnInst *NewRI;
548+
if (RetVal != nullptr)
549+
NewRI = Builder.CreateRet(RetVal->Val);
550+
else
551+
NewRI = Builder.CreateRetVoid();
552+
return Ctx.createReturnInst(NewRI);
553+
}
554+
555+
ReturnInst *ReturnInst::create(Value *RetVal, Instruction *InsertBefore,
556+
Context &Ctx) {
557+
llvm::Instruction *BeforeIR = InsertBefore->getTopmostLLVMInstruction();
558+
auto &Builder = Ctx.getLLVMIRBuilder();
559+
Builder.SetInsertPoint(BeforeIR);
560+
return createCommon(RetVal, Builder, Ctx);
561+
}
562+
563+
ReturnInst *ReturnInst::create(Value *RetVal, BasicBlock *InsertAtEnd,
564+
Context &Ctx) {
565+
auto &Builder = Ctx.getLLVMIRBuilder();
566+
Builder.SetInsertPoint(cast<llvm::BasicBlock>(InsertAtEnd->Val));
567+
return createCommon(RetVal, Builder, Ctx);
568+
}
569+
570+
Value *ReturnInst::getReturnValue() const {
571+
auto *LLVMRetVal = cast<llvm::ReturnInst>(Val)->getReturnValue();
572+
return LLVMRetVal != nullptr ? Ctx.getValue(LLVMRetVal) : nullptr;
573+
}
574+
575+
#ifndef NDEBUG
576+
void ReturnInst::dump(raw_ostream &OS) const {
577+
dumpCommonPrefix(OS);
578+
dumpCommonSuffix(OS);
579+
}
580+
581+
void ReturnInst::dump() const {
582+
dump(dbgs());
583+
dbgs() << "\n";
584+
}
543585

544586
void OpaqueInst::dump(raw_ostream &OS) const {
545587
dumpCommonPrefix(OS);
@@ -626,7 +668,7 @@ Value *Context::registerValue(std::unique_ptr<Value> &&VPtr) {
626668
"Can't register a user!");
627669
Value *V = VPtr.get();
628670
[[maybe_unused]] auto Pair =
629-
LLVMValueToValueMap.insert({VPtr->Val, std::move(VPtr)});
671+
LLVMValueToValueMap.insert({VPtr->Val, std::move(VPtr)});
630672
assert(Pair.second && "Already exists!");
631673
return V;
632674
}
@@ -668,6 +710,11 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
668710
It->second = std::unique_ptr<StoreInst>(new StoreInst(LLVMSt, *this));
669711
return It->second.get();
670712
}
713+
case llvm::Instruction::Ret: {
714+
auto *LLVMRet = cast<llvm::ReturnInst>(LLVMV);
715+
It->second = std::unique_ptr<ReturnInst>(new ReturnInst(LLVMRet, *this));
716+
return It->second.get();
717+
}
671718
default:
672719
break;
673720
}
@@ -696,6 +743,11 @@ StoreInst *Context::createStoreInst(llvm::StoreInst *SI) {
696743
return cast<StoreInst>(registerValue(std::move(NewPtr)));
697744
}
698745

746+
ReturnInst *Context::createReturnInst(llvm::ReturnInst *I) {
747+
auto NewPtr = std::unique_ptr<ReturnInst>(new ReturnInst(I, *this));
748+
return cast<ReturnInst>(registerValue(std::move(NewPtr)));
749+
}
750+
699751
Value *Context::getValue(llvm::Value *V) const {
700752
auto It = LLVMValueToValueMap.find(V);
701753
if (It != LLVMValueToValueMap.end())

‎llvm/unittests/SandboxIR/SandboxIRTest.cpp

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ define i32 @foo(i32 %v0, i32 %v1) {
132132
auto *Arg1 = F.getArg(1);
133133
auto It = BB.begin();
134134
auto *I0 = &*It++;
135-
auto *Ret = &*It++;
135+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
136136

137137
SmallVector<sandboxir::Argument *> Args{Arg0, Arg1};
138138
unsigned OpIdx = 0;
@@ -245,7 +245,7 @@ define i32 @foo(i32 %arg0, i32 %arg1) {
245245
auto *I0 = &*It++;
246246
auto *I1 = &*It++;
247247
auto *I2 = &*It++;
248-
auto *Ret = &*It++;
248+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
249249

250250
bool Replaced;
251251
// Try to replace an operand that doesn't match.
@@ -401,7 +401,7 @@ void @foo(i32 %arg0, i32 %arg1) {
401401
br label %bb1 ; SB3. (Opaque)
402402
403403
bb1:
404-
ret void ; SB5. (Opaque)
404+
ret void ; SB5. (Ret)
405405
}
406406
)IR");
407407
}
@@ -488,7 +488,7 @@ define void @foo(i8 %v1) {
488488
auto It = BB->begin();
489489
auto *I0 = &*It++;
490490
auto *I1 = &*It++;
491-
auto *Ret = &*It++;
491+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
492492

493493
// Check getPrevNode().
494494
EXPECT_EQ(Ret->getPrevNode(), I1);
@@ -508,7 +508,7 @@ define void @foo(i8 %v1) {
508508
// Check getOpcode().
509509
EXPECT_EQ(I0->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
510510
EXPECT_EQ(I1->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
511-
EXPECT_EQ(Ret->getOpcode(), sandboxir::Instruction::Opcode::Opaque);
511+
EXPECT_EQ(Ret->getOpcode(), sandboxir::Instruction::Opcode::Ret);
512512

513513
// Check moveBefore(I).
514514
I1->moveBefore(I0);
@@ -576,7 +576,7 @@ define void @foo(ptr %arg0, ptr %arg1) {
576576
auto *BB = &*F->begin();
577577
auto It = BB->begin();
578578
auto *Ld = cast<sandboxir::LoadInst>(&*It++);
579-
auto *Ret = &*It++;
579+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
580580

581581
// Check getPointerOperand()
582582
EXPECT_EQ(Ld->getPointerOperand(), Arg0);
@@ -607,7 +607,7 @@ define void @foo(i8 %val, ptr %ptr) {
607607
auto *BB = &*F->begin();
608608
auto It = BB->begin();
609609
auto *St = cast<sandboxir::StoreInst>(&*It++);
610-
auto *Ret = &*It++;
610+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
611611

612612
// Check that the StoreInst has been created correctly.
613613
// Check getPointerOperand()
@@ -624,3 +624,42 @@ define void @foo(i8 %val, ptr %ptr) {
624624
EXPECT_EQ(NewSt->getPointerOperand(), Ptr);
625625
EXPECT_EQ(NewSt->getAlign(), 8);
626626
}
627+
628+
TEST_F(SandboxIRTest, ReturnInst) {
629+
parseIR(C, R"IR(
630+
define i8 @foo(i8 %val) {
631+
%add = add i8 %val, 42
632+
ret i8 %val
633+
}
634+
)IR");
635+
llvm::Function *LLVMF = &*M->getFunction("foo");
636+
sandboxir::Context Ctx(C);
637+
sandboxir::Function *F = Ctx.createFunction(LLVMF);
638+
auto *Val = F->getArg(0);
639+
auto *BB = &*F->begin();
640+
auto It = BB->begin();
641+
It++;
642+
auto *Ret = cast<sandboxir::ReturnInst>(&*It++);
643+
644+
// Check that the ReturnInst has been created correctly.
645+
// Check getReturnValue().
646+
EXPECT_EQ(Ret->getReturnValue(), Val);
647+
648+
// Check create(InsertBefore) a void ReturnInst.
649+
auto *NewRet1 = cast<sandboxir::ReturnInst>(
650+
sandboxir::ReturnInst::create(nullptr, /*InsertBefore=*/Ret, Ctx));
651+
EXPECT_EQ(NewRet1->getReturnValue(), nullptr);
652+
// Check create(InsertBefore) a non-void ReturnInst.
653+
auto *NewRet2 = cast<sandboxir::ReturnInst>(
654+
sandboxir::ReturnInst::create(Val, /*InsertBefore=*/Ret, Ctx));
655+
EXPECT_EQ(NewRet2->getReturnValue(), Val);
656+
657+
// Check create(InsertAtEnd) a void ReturnInst.
658+
auto *NewRet3 = cast<sandboxir::ReturnInst>(
659+
sandboxir::ReturnInst::create(nullptr, /*InsertAtEnd=*/BB, Ctx));
660+
EXPECT_EQ(NewRet3->getReturnValue(), nullptr);
661+
// Check create(InsertAtEnd) a non-void ReturnInst.
662+
auto *NewRet4 = cast<sandboxir::ReturnInst>(
663+
sandboxir::ReturnInst::create(Val, /*InsertAtEnd=*/BB, Ctx));
664+
EXPECT_EQ(NewRet4->getReturnValue(), Val);
665+
}

0 commit comments

Comments
 (0)