Remove the OpKernel::AsAsync() const overload.

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

View File

@ -1756,7 +1756,7 @@ void OpKernelContext::CtxFailureWithWarning(const char* file, int line,
void CheckNotInComputeAsync(OpKernelContext* ctx,
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.";
}

View File

@ -136,7 +136,6 @@ class OpKernel {
// Returns nullptr iff this op kernel is synchronous.
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
// 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;
AsyncOpKernel* AsAsync() override { return this; }
const AsyncOpKernel* AsAsync() const override { return this; }
void Compute(OpKernelContext* context) override;
};
@ -1374,6 +1372,10 @@ class OpKernelContext {
};
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);
};

View File

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

View File

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