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:
parent
a007002f21
commit
a7ede3b86c
@ -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.";
|
||||
}
|
||||
|
||||
|
@ -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);
|
||||
};
|
||||
|
||||
|
@ -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(
|
||||
|
@ -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;
|
||||
|
Loading…
Reference in New Issue
Block a user