experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Remove the OpKernel::AsAsync() const overload.

Browse Source 
This overload was only used in `CheckNotInComputeAsync(OpKernelContext*, ...)`, which is modified to have friend access to a non-const `OpKernel*` instead. Removing the overload avoids the need to specify the same logic twice, and the risk of implementing one but not the other (a mistake I made on a recent change).

PiperOrigin-RevId: 298691746
Change-Id: Iba262216fdb6bdce4171794b08acbee393ac116b
This commit is contained in:
Derek Murray 2020-03-03 14:14:47 -08:00 committed by TensorFlower Gardener
parent a007002f21
commit a7ede3b86c
4 changed files with 5 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/framework/op_kernel.cc
View File

@ -1756,7 +1756,7 @@ void OpKernelContext::CtxFailureWithWarning(const char* file, int line,
void CheckNotInComputeAsync(OpKernelContext* ctx, void CheckNotInComputeAsync(OpKernelContext* ctx,
const char* correct_macro_name) { const char* correct_macro_name) {
CHECK_EQ(nullptr, ctx->op_kernel().AsAsync()) CHECK_EQ(nullptr, ctx->params_->op_kernel->AsAsync())
<< "Use " << correct_macro_name << " in AsyncOpKernel implementations."; << "Use " << correct_macro_name << " in AsyncOpKernel implementations.";
} }

6
tensorflow/core/framework/op_kernel.h
View File

@ -136,7 +136,6 @@ class OpKernel {
// Returns nullptr iff this op kernel is synchronous. // Returns nullptr iff this op kernel is synchronous.
virtual AsyncOpKernel* AsAsync() { return nullptr; } virtual AsyncOpKernel* AsAsync() { return nullptr; }
virtual const AsyncOpKernel* AsAsync() const { return nullptr; }
// Initial time (in CPU cycles) we expect an operation to take. Used to // Initial time (in CPU cycles) we expect an operation to take. Used to
// determine whether an operation should be place in a threadpool. Operations // determine whether an operation should be place in a threadpool. Operations
@ -263,7 +262,6 @@ class AsyncOpKernel : public OpKernel {
virtual void ComputeAsync(OpKernelContext* context, DoneCallback done) = 0; virtual void ComputeAsync(OpKernelContext* context, DoneCallback done) = 0;
AsyncOpKernel* AsAsync() override { return this; } AsyncOpKernel* AsAsync() override { return this; }
const AsyncOpKernel* AsAsync() const override { return this; }
void Compute(OpKernelContext* context) override; void Compute(OpKernelContext* context) override;
}; };
@ -1374,6 +1372,10 @@ class OpKernelContext {
}; };
std::unique_ptr<TrackingState> tracking_state_; std::unique_ptr<TrackingState> tracking_state_;
// For access to `params_->op_kernel`.
friend void CheckNotInComputeAsync(OpKernelContext* ctx,
const char* correct_macro_name);
TF_DISALLOW_COPY_AND_ASSIGN(OpKernelContext); TF_DISALLOW_COPY_AND_ASSIGN(OpKernelContext);
}; };

3
tensorflow/core/kernels/data/iterator_ops.cc
View File

@ -902,9 +902,6 @@ class OneShotIteratorOp : public AsyncOpKernel {
AsyncOpKernel* IteratorGetNextOp::AsAsync() { AsyncOpKernel* IteratorGetNextOp::AsAsync() {
return type_string() == "IteratorGetNextSync" ? nullptr : this; return type_string() == "IteratorGetNextSync" ? nullptr : this;
} }
const AsyncOpKernel* IteratorGetNextOp::AsAsync() const {
return type_string() == "IteratorGetNextSync" ? nullptr : this;
}
Status IteratorGetNextOp::DoCompute(OpKernelContext* ctx) { Status IteratorGetNextOp::DoCompute(OpKernelContext* ctx) {
profiler::TraceMe traceme( profiler::TraceMe traceme(

1
tensorflow/core/kernels/data/iterator_ops.h
View File

@ -205,7 +205,6 @@ class IteratorGetNextOp : public HybridAsyncOpKernel {
: HybridAsyncOpKernel(ctx, "tf_data_iterator_get_next") {} : HybridAsyncOpKernel(ctx, "tf_data_iterator_get_next") {}
AsyncOpKernel* AsAsync() override; AsyncOpKernel* AsAsync() override;
const AsyncOpKernel* AsAsync() const override;
protected: protected:
Status DoCompute(OpKernelContext* ctx) override; Status DoCompute(OpKernelContext* ctx) override;