Add target into the record instruction

rir/src/compiler/native/builtins.cpp

@@ -957,42 +957,30 @@ void deoptImpl(rir::Code* c, SEXP cls, DeoptMetadata* m, R_bcstack_t* args,
     assert(false);
 }
 
-void recordTypefeedbackImpl(rir::Code* code, uint32_t idx, SEXP value) {
+void recordTypefeedbackImpl(rir::TypeFeedback* typeFeedback, uint32_t idx,
+                            SEXP value) {
     // we cannot pass the feedback directly because the call to this builtin is
     // generated from places that do not have access to the feedback vector
-    auto dt = code->function()->dispatchTable();
-    auto baseline = dt->baseline();
-    auto& slot = baseline->typeFeedback()[idx];
-
-    switch (slot.kind) {
-    case TypeFeedbackKind::Call: {
-        auto& feedback = slot.callees();
-        feedback.record(baseline->body(), value);
-        break;
-    }
-    case TypeFeedbackKind::Test: {
-        auto& feedback = slot.test();
-        feedback.record(value);
-        break;
-    }
-    case TypeFeedbackKind::Type: {
+    typeFeedback->record(idx, value);
+
+    auto& slot = (*typeFeedback)[idx];
+    if (slot.kind == TypeFeedbackKind::Type) {
         auto& feedback = slot.type();
-        feedback.record(value);
 
         if (TYPEOF(value) == PROMSXP) {
             if (PRVALUE(value) == R_UnboundValue &&
-                feedback.stateBeforeLastForce < ObservedValues::promise)
+                feedback.stateBeforeLastForce < ObservedValues::promise) {
                 feedback.stateBeforeLastForce = ObservedValues::promise;
-            else if (feedback.stateBeforeLastForce <
-                     ObservedValues::evaluatedPromise)
+            } else if (feedback.stateBeforeLastForce <
+                       ObservedValues::evaluatedPromise) {
                 feedback.stateBeforeLastForce =
                     ObservedValues::evaluatedPromise;
+            }
         } else {
             if (feedback.stateBeforeLastForce < ObservedValues::value)
                 feedback.stateBeforeLastForce = ObservedValues::value;
         }
     }
-    }
 }
 
 void assertFailImpl(const char* msg) {

rir/src/compiler/native/lower_function_llvm.cpp

@@ -3602,7 +3602,8 @@ void LowerFunctionLLVM::compile() {
 
                 call(
                     NativeBuiltins::get(NativeBuiltins::Id::recordTypefeedback),
-                    {paramCode(), c(rec->idx), loadSxp(rec->arg(0).val())});
+                    {convertToPointer(rec->feedback, t::i8, true), c(rec->idx),
+                     loadSxp(rec->arg(0).val())});
 
                 break;
             }
@@ -6132,19 +6133,28 @@ void LowerFunctionLLVM::compile() {
 
             // For OSR-in try to collect more typefeedback for the part of the
             // code that was not yet executed.
+            // FIXME: is this correct? the feedbackOrigin index?
             if (cls->isContinuation() && Rep::Of(i) == Rep::SEXP &&
                 variables_.count(i) &&
                 !cls->isContinuation()->continuationContext->asDeoptContext()) {
                 if (i->hasTypeFeedback()) {
-                    call(NativeBuiltins::get(
-                             NativeBuiltins::Id::recordTypefeedback),
-                         {paramCode(),
-                          c(i->typeFeedback().feedbackOrigin.idx()), load(i)});
+                    auto& origin = i->typeFeedback().feedbackOrigin;
+                    if (origin.isValid()) {
+                        call(NativeBuiltins::get(
+                                 NativeBuiltins::Id::recordTypefeedback),
+                             {convertToPointer(
+                                  &origin.function()->typeFeedback(), t::i8,
+                                  true),
+                              c(origin.idx()), load(i)});
+                    }
                 }
                 if (i->hasCallFeedback()) {
                     call(NativeBuiltins::get(
                              NativeBuiltins::Id::recordTypefeedback),
-                         {paramCode(),
+                         {convertToPointer(&i->callFeedback()
+                                                .feedbackOrigin.function()
+                                                ->typeFeedback(),
+                                           t::i8, true),
                           c(i->typeFeedback().feedbackOrigin.idx()), load(i)});
                 }
             }

rir/src/compiler/pir/builder.cpp

@@ -126,6 +126,7 @@ Builder::Builder(Continuation* cnt, Value* closureEnv)
         }
         auto mkenv = new MkEnv(closureEnv, names, args.data(), miss);
 
+        // FIXME: what does this mean, we need both rirFun and we need idx
         mkenv->updateTypeFeedback().feedbackOrigin.function(
             cnt->owner()->rirFunction());
         add(mkenv);
@@ -173,6 +174,7 @@ Builder::Builder(ClosureVersion* version, Value* closureEnv)
     auto rirFun = version->owner()->rirFunction();
     if (rirFun->flags.contains(rir::Function::NeedsFullEnv))
         mkenv->neverStub = true;
+    // FIXME: what does this mean, we need both rirFun and we need idx
     mkenv->updateTypeFeedback().feedbackOrigin.function(rirFun);
     add(mkenv);
     this->env = mkenv;

rir/src/compiler/pir/instruction.cpp

@@ -1019,10 +1019,11 @@ bool Deopt::hasDeoptReason() const {
     return deoptReason() != DeoptReasonWrapper::unknown();
 }
 
-Record::Record(rir::TypeFeedbackKind kind, uint32_t idx)
+Record::Record(rir::TypeFeedback* feedback, rir::TypeFeedbackKind kind,
+               uint32_t idx)
     : FixedLenInstruction(PirType::voyd(), {{PirType::any()}},
                           {{UnknownDeoptTrigger::instance()}}),
-      kind(kind), idx(idx) {}
+      feedback(feedback), kind(kind), idx(idx) {}
 
 Value* Record::getValue() const { return arg<0>().val(); }
 void Record::setValue(Value* value) { arg<0>().val() = value; }
@@ -1040,7 +1041,7 @@ void Record::printArgs(std::ostream& out, bool tty) const {
         break;
     }
 
-    out << "#" << idx << " ";
+    out << "#" << idx << " (" << feedback << ") ";
 
     getValue()->printRef(out);
 }

rir/src/compiler/pir/instruction.h

@@ -2720,10 +2720,12 @@ class Record
     : public FixedLenInstruction<Tag::Record, Record, 1, Effects::AnyI(),
                                  HasEnvSlot::No, Controlflow::None> {
   public:
+    rir::TypeFeedback* feedback;
     rir::TypeFeedbackKind kind;
     uint32_t idx;
 
-    explicit Record(rir::TypeFeedbackKind kind, unsigned idx);
+    explicit Record(rir::TypeFeedback* feedback, rir::TypeFeedbackKind kind,
+                    unsigned idx);
     Value* getValue() const;
     void setValue(Value* callee);
     void printArgs(std::ostream& out, bool tty) const override;

rir/src/compiler/rir2pir/rir2pir.cpp

@@ -378,7 +378,9 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
         Value* target = top();
 
         if (baseline) {
-            auto rec = insert(new Record(rir::TypeFeedbackKind::Test, idx));
+            auto feedback = &srcCode->function()->typeFeedback();
+            auto rec =
+                insert(new Record(feedback, rir::TypeFeedbackKind::Test, idx));
             rec->setValue(target);
         } else {
             auto& feedback = typeFeedback.test(idx);
@@ -413,7 +415,9 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
         Value* target = top();
 
         if (baseline) {
-            auto rec = insert(new Record(rir::TypeFeedbackKind::Type, idx));
+            auto feedback = &srcCode->function()->typeFeedback();
+            auto rec =
+                insert(new Record(feedback, rir::TypeFeedbackKind::Type, idx));
             rec->setValue(target);
         } else {
             auto& feedback = typeFeedback.types(idx);
@@ -457,7 +461,9 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
         Value* target = top();
 
         if (baseline) {
-            auto rec = insert(new Record(rir::TypeFeedbackKind::Call, idx));
+            auto feedback = &srcCode->function()->typeFeedback();
+            auto rec =
+                insert(new Record(feedback, rir::TypeFeedbackKind::Call, idx));
             rec->setValue(target);
         } else {
             const auto& feedback = typeFeedback.callees(bc.immediate.i);
@@ -490,7 +496,7 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
                 if (feedback.numTargets == 1) {
                     assert(!feedback.invalid &&
                            "feedback can't be invalid if numTargets is 1");
-                    f.monomorphic = feedback.getTarget(srcCode, 0);
+                    f.monomorphic = feedback.getTarget(srcCode->function(), 0);
                     f.type = TYPEOF(f.monomorphic);
                     f.stableEnv = true;
                 } else if (feedback.numTargets > 1) {
@@ -499,7 +505,7 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
                     bool stableBody = !feedback.invalid;
                     bool stableEnv = !feedback.invalid;
                     for (size_t i = 0; i < feedback.numTargets; ++i) {
-                        SEXP b = feedback.getTarget(srcCode, i);
+                        SEXP b = feedback.getTarget(srcCode->function(), i);
                         if (!first) {
                             first = b;
                         } else {
@@ -524,7 +530,8 @@ bool Rir2Pir::compileBC(const BC& bc, Opcode* pos, Opcode* nextPos,
                                         d->callTargetTrigger();
                                     for (size_t i = 0; i < feedback.numTargets;
                                          ++i) {
-                                        SEXP b = feedback.getTarget(srcCode, i);
+                                        SEXP b = feedback.getTarget(
+                                            srcCode->function(), i);
                                         if (b != deoptCallTarget)
                                             deoptedCallTargets.insert(b);
                                     }

rir/src/runtime/TypeFeedback.cpp

@@ -32,9 +32,9 @@ void ObservedCallees::record(Code* caller, SEXP callee,
     }
 }
 
-SEXP ObservedCallees::getTarget(const Code* code, size_t pos) const {
+SEXP ObservedCallees::getTarget(const Function* function, size_t pos) const {
     assert(pos < numTargets);
-    return code->getExtraPoolEntry(targets[pos]);
+    return function->body()->getExtraPoolEntry(targets[pos]);
 }
 
 FeedbackOrigin::FeedbackOrigin(rir::Function* function, uint32_t idx)
@@ -87,7 +87,7 @@ void DeoptReason::record(SEXP val) const {
     }
 }
 
-void ObservedCallees::print(std::ostream& out, const Code* code) const {
+void ObservedCallees::print(std::ostream& out, const Function* function) const {
     out << "callees: ";
     if (taken == ObservedCallees::CounterOverflow)
         out << "*, <";
@@ -102,7 +102,7 @@ void ObservedCallees::print(std::ostream& out, const Code* code) const {
     out << (numTargets ? ", " : " ");
 
     for (unsigned i = 0; i < numTargets; ++i) {
-        auto target = getTarget(code, i);
+        auto target = getTarget(function, i);
         out << target << "(" << Rf_type2char(TYPEOF(target)) << ") ";
     }
 }
@@ -170,7 +170,7 @@ void TypeFeedbackSlot::print(std::ostream& out,
                              const Function* function) const {
     switch (kind) {
     case TypeFeedbackKind::Call:
-        feedback_.callees.print(out, function->body());
+        feedback_.callees.print(out, function);
         break;
     case TypeFeedbackKind::Test:
         feedback_.test.print(out);
@@ -182,7 +182,8 @@ void TypeFeedbackSlot::print(std::ostream& out,
 }
 
 void TypeFeedback::print(std::ostream& out) const {
-    std::cout << "== type feedback ==" << std::endl;
+    std::cout << "== type feedback " << this << " (fun " << owner_
+              << ") ==" << std::endl;
     int i = 0;
     for (auto& slot : slots_) {
         out << "#" << i++ << ": ";
@@ -214,4 +215,7 @@ TypeFeedbackSlot* FeedbackOrigin::slot() const {
         return nullptr;
     }
 }
+bool FeedbackOrigin::isValid() const {
+    return function_ != nullptr && function_->typeFeedback().size() > idx_;
+}
 } // namespace rir

rir/src/runtime/TypeFeedback.h

@@ -38,8 +38,8 @@ struct ObservedCallees {
     std::array<unsigned, MaxTargets> targets;
 
     void record(Code* caller, SEXP callee, bool invalidateWhenFull = false);
-    SEXP getTarget(const Code* code, size_t pos) const;
-    void print(std::ostream& out, const Code* code) const;
+    SEXP getTarget(const Function* function, size_t pos) const;
+    void print(std::ostream& out, const Function* function) const;
 };
 
 static_assert(sizeof(ObservedCallees) == 4 * sizeof(uint32_t),
@@ -154,7 +154,7 @@ struct FeedbackOrigin {
     FeedbackOrigin() {}
     FeedbackOrigin(rir::Function* fun, uint32_t idx);
 
-    bool isValid() const { return function_ != nullptr; }
+    bool isValid() const;
     TypeFeedbackSlot* slot() const;
     uint32_t idx() const { return idx_; }
     Function* function() const { return function_; }
@@ -324,6 +324,8 @@ class TypeFeedback {
     void print(std::ostream& out) const;
 
     void record(uint32_t idx, SEXP callee);
+
+    uint32_t size() const { return slots_.size(); }
 };
 
 #pragma pack(pop)