Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit 55e12f7

Browse files
mtrofinmtrofin
committed
[NFC][MLGO] Just use the underlying protobuf object for logging
Avoid buffering just to copy the buffered data, in 'development mode', when logging. Instead, just populate the underlying protobuf. Differential Revision: https://reviews.llvm.org/D106592
1 parent deebf18 commit 55e12f7

File tree

4 files changed

+231
-132
lines changed

4 files changed

+231
-132
lines changed

‎llvm/include/llvm/Analysis/Utils/TFUtils.h

Lines changed: 39 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include "llvm/Config/llvm-config.h"
1313

1414
#ifdef LLVM_HAVE_TF_API
15+
#include "llvm/ADT/StringMap.h"
1516
#include "llvm/IR/LLVMContext.h"
1617
#include "llvm/Support/JSON.h"
1718

@@ -120,56 +121,62 @@ loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
120121
/// The assumption is that, for an event to be logged (i.e. a set of feature
121122
/// values and a reward), the user calls the log* API for each feature exactly
122123
/// once, providing the index matching the position in the feature spec list
123-
/// provided at construction:
124+
/// provided at construction. The example assumes the first feature's element
125+
/// type is float, the second is int64, and the reward is float:
126+
///
124127
/// event 0:
125-
/// logTensorValue(0, ...)
126-
/// logTensorValue(1, ...)
128+
/// logFloatValue(0, ...)
129+
/// logInt64Value(1, ...)
127130
/// ...
128-
/// logReward(...)
131+
/// logFloatReward(...)
129132
/// event 1:
130-
/// logTensorValue(0, ...)
131-
/// logTensorValue(1, ...)
133+
/// logFloatValue(0, ...)
134+
/// logInt64Value(1, ...)
132135
/// ...
133-
/// logReward(...)
136+
/// logFloatReward(...)
134137
///
135138
/// At the end, call print to generate the protobuf.
139+
/// Alternatively, don't call logReward at the end of each event, just
140+
/// log{Float|Int32|Int64}FinalReward at the end.
141+
class LoggerDataImpl;
136142
class Logger final {
137143
public:
138-
/// Construct a Logger. If IncludeReward is false, then logReward shouldn't
139-
/// be called, and the reward feature won't be printed out.
144+
/// Construct a Logger. If IncludeReward is false, then logReward or
145+
/// logFinalReward shouldn't be called, and the reward feature won't be
146+
/// printed out.
140147
Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
141-
const TensorSpec &RewardSpec, bool IncludeReward)
142-
: FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
143-
RawLogData(FeatureSpecs.size() + IncludeReward),
144-
IncludeReward(IncludeReward) {}
145-
146-
template <typename T> void logReward(T Value) {
147-
assert(IncludeReward);
148-
logTensorValue(RawLogData.size() - 1, &Value);
149-
}
148+
const TensorSpec &RewardSpec, bool IncludeReward);
150149

151-
template <typename T> void logFinalReward(T Value) {
152-
assert(RawLogData.back().empty());
153-
logReward(Value);
154-
}
150+
~Logger();
155151

156-
template <typename T>
157-
void logTensorValue(size_t FeatureID, const T *Value, size_t Size = 1) {
158-
const char *Start = reinterpret_cast<const char *>(Value);
159-
const char *End = Start + sizeof(T) * Size;
160-
RawLogData[FeatureID].insert(RawLogData[FeatureID].end(), Start, End);
161-
}
152+
void logFloatReward(float Value);
153+
void logInt32Reward(int32_t Value);
154+
void logInt64Reward(int64_t Value);
155+
156+
void logFloatFinalReward(float Value);
157+
void logInt32FinalReward(int32_t Value);
158+
void logInt64FinalReward(int64_t Value);
159+
160+
void logFloatValue(size_t FeatureID, const float *Value);
161+
void logInt32Value(size_t FeatureID, const int32_t *Value);
162+
void logInt64Value(size_t FeatureID, const int64_t *Value);
163+
164+
void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
165+
166+
// Warning! For int32_t, the return is set up for int64_t, so the caller needs
167+
// to piecemeal cast their int32_t values.
168+
// FIXME: let's drop int32_t support. While it's supported by evaluator, it's
169+
// not supported by the tensorflow::SequenceExample proto. For small values,
170+
// we can consider using bytes.
171+
char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
162172

163173
void print(raw_ostream &OS);
164174

165175
private:
166176
std::vector<LoggedFeatureSpec> FeatureSpecs;
167177
TensorSpec RewardSpec;
168-
/// RawData has one entry per feature, plus one more for the reward.
169-
/// Each feature's values are then stored in a vector, in succession.
170-
/// This means the ith event is stored at [*][i]
171-
std::vector<std::vector<char>> RawLogData;
172178
const bool IncludeReward;
179+
std::unique_ptr<LoggerDataImpl> LoggerData;
173180
};
174181

175182
class TFModelEvaluator final {

‎llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -353,24 +353,22 @@ void TrainingLogger::logInlineEvent(const InlineEvent &Event,
353353
size_t CurrentFeature = 0;
354354
for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
355355
int64_t F = ModelRunner.getFeature(CurrentFeature);
356-
L->logTensorValue(CurrentFeature, &F);
356+
L->logInt64Value(CurrentFeature, &F);
357357
}
358358

359359
for (size_t I = 1; I < OutputCount; ++I) {
360360
const auto &Result = *MUTR->lastEvaluationResult();
361-
auto &Spec = MUTR->outputLoggedFeatureSpecs()[I].Spec;
362361
const char *RawData =
363362
reinterpret_cast<const char *>(Result.getUntypedTensorValue(I));
364-
L->logTensorValue(CurrentFeature, RawData,
365-
Spec.getElementCount() * Spec.getElementByteSize());
363+
L->logSpecifiedTensorValue(CurrentFeature, RawData);
366364
++CurrentFeature;
367365
}
368366

369367
assert(CurrentFeature == DefaultDecisionPos);
370-
L->logTensorValue(DefaultDecisionPos, &Event.DefaultDecision);
371-
L->logTensorValue(DecisionPos, &Event.AdvisedDecision);
368+
L->logInt64Value(DefaultDecisionPos, &Event.DefaultDecision);
369+
L->logInt64Value(DecisionPos, &Event.AdvisedDecision);
372370
if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
373-
L->logReward(Event.Reward);
371+
L->logInt64Reward(Event.Reward);
374372

375373
// For debugging / later use
376374
Effects.push_back(Event.Effect);

‎llvm/lib/Analysis/TFUtils.cpp

Lines changed: 137 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232

3333
using namespace llvm;
3434

35+
using google::protobuf::Message;
36+
using google::protobuf::TextFormat;
37+
3538
static cl::opt<bool>
3639
ProtobufTextMode("tfutils-text-log", cl::init(false), cl::Hidden,
3740
cl::desc("Output textual (human-readable) protobuf."));
@@ -70,55 +73,6 @@ TFStatusPtr createTFStatus() {
7073
TFSessionOptionsPtr createTFSessionOptions() {
7174
return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
7275
}
73-
74-
/// Write a list of tensors as a sequence of TensorFlow FeatureList protobufs.
75-
/// The tensors are assumed to be stored contiguously, in row-major format,
76-
/// in the TensorData buffer. Each tensor has the shape given by Spec. The
77-
/// feature name in the output is either the provided LoggingName, if
78-
/// specified, otherwise it's the name of the tensor (as given by Spec).
79-
void writeRawTensorsAsFeatureLists(tensorflow::FeatureLists *FE,
80-
const LoggedFeatureSpec &LoggedSpec,
81-
const char *TensorData, size_t TensorCount,
82-
bool FinalReward = false) {
83-
const auto &Spec = LoggedSpec.Spec;
84-
// The 'Feature' protobuf only has 3 possible fields: float_list,
85-
// int64_list, or bytes_list, so we capture int32 values as int64. We don't
86-
// support any other types.
87-
tensorflow::FeatureList &FL = (*FE->mutable_feature_list())[(
88-
LoggedSpec.LoggingName ? *LoggedSpec.LoggingName : Spec.name())];
89-
90-
const char *CurrentTensor = TensorData;
91-
const size_t TensorByteSize =
92-
Spec.getElementCount() * Spec.getElementByteSize();
93-
const size_t ElemCount = Spec.getElementCount();
94-
for (size_t E = 0; E < TensorCount; ++E) {
95-
const bool ShouldWrite = E + 1 == TensorCount || !FinalReward;
96-
97-
if (Spec.isElementType<int64_t>()) {
98-
auto *MF = FL.add_feature()->mutable_int64_list()->mutable_value();
99-
MF->Resize(ElemCount, 0);
100-
if (ShouldWrite)
101-
memcpy(MF->mutable_data(), CurrentTensor, TensorByteSize);
102-
} else if (Spec.isElementType<int32_t>()) {
103-
auto *MF = FL.add_feature()->mutable_int64_list()->mutable_value();
104-
MF->Resize(ElemCount, 0);
105-
if (ShouldWrite) {
106-
const int32_t *TD = reinterpret_cast<const int32_t *>(CurrentTensor);
107-
for (size_t I = 0; I < ElemCount; ++I)
108-
(*MF)[I] = TD[I];
109-
}
110-
} else if (Spec.isElementType<float>()) {
111-
auto *MF = FL.add_feature()->mutable_float_list()->mutable_value();
112-
MF->Resize(ElemCount, 0.0);
113-
if (ShouldWrite)
114-
memcpy(MF->mutable_data(), CurrentTensor, TensorByteSize);
115-
} else {
116-
llvm_unreachable("Unsupported tensor type.");
117-
}
118-
if (ShouldWrite)
119-
CurrentTensor += TensorByteSize;
120-
}
121-
}
12276
} // namespace
12377

12478
namespace llvm {
@@ -304,6 +258,76 @@ class TFModelEvaluatorImpl {
304258
bool checkReportAndInvalidate(const TF_Output &Output,
305259
const TensorSpec &OutputSpec);
306260
};
261+
262+
class LoggerDataImpl {
263+
const std::vector<LoggedFeatureSpec> LoggedFeatureSpecs;
264+
const TensorSpec RewardSpec;
265+
266+
tensorflow::SequenceExample SE;
267+
std::vector<tensorflow::FeatureList *> FeatureLists;
268+
tensorflow::FeatureList *Reward = nullptr;
269+
270+
public:
271+
LoggerDataImpl(const std::vector<LoggedFeatureSpec> &LoggedSpecs,
272+
const TensorSpec &RewardSpec, bool IncludeReward)
273+
: LoggedFeatureSpecs(LoggedSpecs), RewardSpec(RewardSpec) {
274+
auto *FL = SE.mutable_feature_lists()->mutable_feature_list();
275+
if (IncludeReward)
276+
Reward = &(*FL)[RewardSpec.name()];
277+
// Allocate first the map entries, then capture their address. We will not
278+
// mutate the set of features after this (i.e. the pointers won't dangle).
279+
for (const auto &LFS : LoggedSpecs) {
280+
(*FL)[LFS.LoggingName ? *LFS.LoggingName : LFS.Spec.name()] = {};
281+
}
282+
for (const auto &LFS : LoggedSpecs)
283+
FeatureLists.push_back(
284+
&(*FL)[LFS.LoggingName ? *LFS.LoggingName : LFS.Spec.name()]);
285+
}
286+
287+
void print(raw_ostream &OS) {
288+
std::string OutStr;
289+
if (ProtobufTextMode)
290+
google::protobuf::TextFormat::PrintToString(SE, &OutStr);
291+
else
292+
OutStr = SE.SerializeAsString();
293+
294+
OS << OutStr;
295+
}
296+
297+
char *addNewTensor(size_t FeatureID) {
298+
const auto &Spec = LoggedFeatureSpecs[FeatureID].Spec;
299+
if (Spec.isElementType<float>()) {
300+
auto *RF = FeatureLists[FeatureID]
301+
->add_feature()
302+
->mutable_float_list()
303+
->mutable_value();
304+
RF->Resize(Spec.getElementCount(), 0.0);
305+
return reinterpret_cast<char *>(RF->mutable_data());
306+
} else if (Spec.isElementType<int32_t>() || Spec.isElementType<int64_t>()) {
307+
auto *RF = FeatureLists[FeatureID]
308+
->add_feature()
309+
->mutable_int64_list()
310+
->mutable_value();
311+
RF->Resize(Spec.getElementCount(), 0);
312+
return reinterpret_cast<char *>(RF->mutable_data());
313+
}
314+
llvm_unreachable("Unsupported tensor type.");
315+
}
316+
317+
template <typename T> void logReward(T Value) {
318+
if (RewardSpec.isElementType<float>())
319+
Reward->add_feature()->mutable_float_list()->add_value(Value);
320+
else if (RewardSpec.isElementType<int32_t>() ||
321+
RewardSpec.isElementType<int64_t>())
322+
Reward->add_feature()->mutable_int64_list()->add_value(Value);
323+
else
324+
llvm_unreachable("Unsupported tensor type.");
325+
}
326+
327+
size_t getNrRecords() const {
328+
return FeatureLists.empty() ? 0 : FeatureLists[0]->feature().size();
329+
}
330+
};
307331
} // namespace llvm
308332

309333
TFModelEvaluatorImpl::TFModelEvaluatorImpl(
@@ -448,37 +472,71 @@ TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_IMPL)
448472
TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
449473
TFModelEvaluator::~TFModelEvaluator() {}
450474

451-
void Logger::print(raw_ostream &OS) {
452-
tensorflow::SequenceExample SE;
475+
Logger::Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
476+
const TensorSpec &RewardSpec, bool IncludeReward)
477+
: FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
478+
IncludeReward(IncludeReward),
479+
LoggerData(std::make_unique<LoggerDataImpl>(FeatureSpecs, RewardSpec,
480+
IncludeReward)) {}
453481

454-
if (RawLogData.empty())
455-
return;
456-
if (RawLogData[0].empty())
457-
return;
458-
size_t Tensor0Size = FeatureSpecs[0].Spec.getElementCount() *
459-
FeatureSpecs[0].Spec.getElementByteSize();
460-
size_t NumberOfRecords = RawLogData[0].size() / Tensor0Size;
461-
if (NumberOfRecords == 0)
462-
return;
463-
size_t RewardSize =
464-
RewardSpec.getElementCount() * RewardSpec.getElementByteSize();
465-
size_t NumberOfRewards = RawLogData.back().size() / RewardSize;
466-
467-
tensorflow::FeatureLists *FE = SE.mutable_feature_lists();
468-
for (size_t I = 0; I < FeatureSpecs.size(); ++I)
469-
writeRawTensorsAsFeatureLists(FE, FeatureSpecs[I], RawLogData[I].data(),
470-
NumberOfRecords);
471-
472-
if (IncludeReward)
473-
writeRawTensorsAsFeatureLists(FE, {RewardSpec, None},
474-
RawLogData.back().data(), NumberOfRecords,
475-
NumberOfRewards == 1);
476-
std::string OutStr;
477-
if (ProtobufTextMode) {
478-
google::protobuf::TextFormat::PrintToString(SE, &OutStr);
479-
} else {
480-
OutStr = SE.SerializeAsString();
482+
Logger::~Logger() {}
483+
484+
#define LOG_REWARD(NAME, TYPE) \
485+
void Logger::log##NAME##Reward(TYPE Value) { \
486+
assert(IncludeReward); \
487+
LoggerData->logReward(Value); \
481488
}
482-
OS << OutStr;
489+
490+
LOG_REWARD(Float, float)
491+
LOG_REWARD(Int32, int32_t)
492+
LOG_REWARD(Int64, int64_t)
493+
#undef LOG_REWARD
494+
495+
#define LOG_FINAL_REWARD(NAME, TYPE) \
496+
void Logger::log##NAME##FinalReward(TYPE Value) { \
497+
assert(RewardSpec.isElementType<TYPE>()); \
498+
for (size_t I = 1; I < LoggerData->getNrRecords(); ++I) \
499+
log##NAME##Reward(0); \
500+
log##NAME##Reward(Value); \
501+
}
502+
503+
LOG_FINAL_REWARD(Float, float)
504+
LOG_FINAL_REWARD(Int32, int32_t)
505+
LOG_FINAL_REWARD(Int64, int64_t)
506+
#undef LOG_FINAL_REWARD
507+
508+
void Logger::logFloatValue(size_t FeatureID, const float *Value) {
509+
assert(FeatureSpecs[FeatureID].Spec.isElementType<float>());
510+
logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
511+
}
512+
513+
void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
514+
assert(FeatureSpecs[FeatureID].Spec.isElementType<int64_t>());
515+
logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
516+
}
517+
518+
void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
519+
assert(FeatureSpecs[FeatureID].Spec.isElementType<int32_t>());
520+
logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
483521
}
522+
523+
void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
524+
const auto &Spec = FeatureSpecs[FeatureID].Spec;
525+
char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
526+
if (Spec.isElementType<int32_t>())
527+
for (size_t I = 0; I < Spec.getElementCount(); ++I)
528+
(reinterpret_cast<int64_t *>(Buff))[I] =
529+
static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
530+
else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
531+
std::memcpy(Buff, RawData,
532+
Spec.getElementCount() * Spec.getElementByteSize());
533+
else
534+
llvm_unreachable("Unsupported tensor type");
535+
}
536+
537+
char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
538+
return reinterpret_cast<char *>(LoggerData->addNewTensor(FeatureID));
539+
}
540+
541+
void Logger::print(raw_ostream &OS) { LoggerData->print(OS); }
484542
#endif // defined(LLVM_HAVE_TF_API)

0 commit comments

Comments
 (0)