diff --git a/tf_adapter/kernels/aicore/npu_mixed_precesion_ops.cc b/tf_adapter/kernels/aicore/npu_mixed_precesion_ops.cc index 753d608a6941ec34742cf6bc1f63069a339f2be3..c86b4d816b3e7e034d52d0822a9a5e099f851f7e 100644 --- a/tf_adapter/kernels/aicore/npu_mixed_precesion_ops.cc +++ b/tf_adapter/kernels/aicore/npu_mixed_precesion_ops.cc @@ -99,4 +99,29 @@ class NpuClearFloatStatusV2Op : public OpKernel { }; REGISTER_KERNEL_BUILDER(Name("NpuClearFloatStatusV2").Device(DEVICE_CPU), NpuClearFloatStatusV2Op); + +class NpuClearFloatDebugStatusOp : public OpKernel { + public: + explicit NpuClearFloatDebugStatusOp(OpKernelConstruction *context) : OpKernel(context) {} + ~NpuClearFloatDebugStatusOp() override = default; + void Compute(OpKernelContext *context) override { + (void) context; + LOG(INFO) << "NpuClearFloatDebugStatus Compute"; + } +}; + +REGISTER_KERNEL_BUILDER(Name("NpuClearFloatDebugStatus").Device(DEVICE_CPU), NpuClearFloatDebugStatusOp); + +class NpuGetFloatDebugStatusOp : public OpKernel { + public: + explicit NpuGetFloatDebugStatusOp(OpKernelConstruction *context) : OpKernel(context) {} + ~NpuGetFloatDebugStatusOp() override = default; + void Compute(OpKernelContext *context) override { + (void) context; + LOG(INFO) << "NpuGetFloatDebugStatus Compute"; + } +}; + +REGISTER_KERNEL_BUILDER(Name("NpuGetFloatDebugStatus").Device(DEVICE_CPU), NpuGetFloatDebugStatusOp); + } // namespace tensorflow diff --git a/tf_adapter/ops/aicore/npu_mixed_precesion_ops.cc b/tf_adapter/ops/aicore/npu_mixed_precesion_ops.cc index 8a1af5d5074cf4600dfa12df12842bcb05f274db..16133576e2c2530f34494edd427d865482fcb625 100644 --- a/tf_adapter/ops/aicore/npu_mixed_precesion_ops.cc +++ b/tf_adapter/ops/aicore/npu_mixed_precesion_ops.cc @@ -65,6 +65,16 @@ REGISTER_OP("NpuGetFloatStatusV2") return Status::OK(); }); +REGISTER_OP("NpuGetFloatDebugStatus") + .Output("data: int32") + .SetIsStateful() + .SetShapeFn([](InferenceContext *c) { + std::vector output_dims; + output_dims.emplace_back(c->MakeDim(8)); + auto output_shape = c->MakeShape(output_dims); + c->set_output(0, output_shape); + return Status::OK(); + }); REGISTER_OP("NpuClearFloatStatus") .Input("float_status: float") @@ -86,4 +96,10 @@ REGISTER_OP("NpuClearFloatStatusV2") Set the value of global workspace to 0. )doc") .SetIsStateful(); + +REGISTER_OP("NpuClearFloatDebugStatus") + .Doc(R"doc( + Set the value of global workspace to 0. + )doc") + .SetIsStateful(); } // namespace tensorflow