Opsapi add layernormgrad

musa_ext/kernels/fusion/musa_gelu_grad_op.cc

@@ -0,0 +1,76 @@
+#include "../utils_op.h"
+#include "tensorflow/core/framework/bfloat16.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+namespace musa {
+
+template <typename T>
+class MusaGeluGradOp : public MusaOpKernel {
+ public:
+  explicit MusaGeluGradOp(OpKernelConstruction* ctx) : MusaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("approximate", &approximate_));
+  }
+
+  bool IsExpensive() override { return false; }
+
+  void Compute(OpKernelContext* ctx) override {
+    const Tensor& dy = ctx->input(0);
+    const Tensor& x = ctx->input(1);
+
+    OP_REQUIRES(ctx, dy.shape() == x.shape(),
+                errors::InvalidArgument(
+                    "dy and x must have the same shape. dy: ",
+                    dy.shape().DebugString(), ", x: ", x.shape().DebugString()));
+
+    Tensor* dx = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &dx));
+
+    if (x.NumElements() == 0) return;
+
+    auto& handle = GetHandleByCtx(ctx);
+    auto mt_dy = CreateMTensor(dy, format_);
+    auto mt_x = CreateMTensor(x, format_);
+    auto mt_dx = CreateMTensor(*dx, format_);
+
+    ::musa::dnn::Binary op;
+    const BINARY_MODE mode = approximate_ ? BINARY_MODE::GELU_TANH_BW
+                                          : BINARY_MODE::GELU_NONE_BW;
+    MTOP_CHECK_OK(op.SetMode(mode), "Set GELU_BW Mode", ctx);
+    MTOP_CHECK_OK_RUN(op.Run(handle, mt_dx, mt_dy, mt_x), "GELU_BW Run", ctx);
+  }
+
+ private:
+  bool approximate_;
+};
+
+#define REGISTER_MUSA_GELU_GRAD(TYPE)                                \
+  REGISTER_KERNEL_BUILDER(                                           \
+      Name("MusaGeluGrad").Device("MUSA").TypeConstraint<TYPE>("T"), \
+      MusaGeluGradOp<TYPE>);
+
+REGISTER_MUSA_GELU_GRAD(float);
+REGISTER_MUSA_GELU_GRAD(double);
+REGISTER_MUSA_GELU_GRAD(Eigen::half);
+REGISTER_MUSA_GELU_GRAD(bfloat16);
+
+#undef REGISTER_MUSA_GELU_GRAD
+
+}  // namespace musa
+
+REGISTER_OP("MusaGeluGrad")
+    .Input("dy: T")
+    .Input("x: T")
+    .Output("dx: T")
+    .Attr("T: {float, double, half, bfloat16}")
+    .Attr("approximate: bool = false")
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(1));
+      return OkStatus();
+    });
+
+}  // namespace tensorflow

musa_ext/kernels/fusion/musa_layernorm_grad_op.cc

@@ -0,0 +1,160 @@
+#include <vector>
+
+#include "../utils_op.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace tensorflow {
+namespace musa {
+
+template <typename T>
+class MusaLayerNormGradOp : public MusaOpKernel {
+ public:
+  explicit MusaLayerNormGradOp(OpKernelConstruction* ctx) : MusaOpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_));
+  }
+
+  bool IsExpensive() override { return true; }
+
+  void Compute(OpKernelContext* ctx) override {
+    const Tensor& dy = ctx->input(0);
+    const Tensor& x = ctx->input(1);
+    const Tensor& gamma = ctx->input(2);
+    const Tensor& beta = ctx->input(3);
+
+    OP_REQUIRES(ctx, x.dims() >= 1,
+                errors::InvalidArgument("Input rank must be >= 1"));
+    OP_REQUIRES(ctx, dy.shape() == x.shape(),
+                errors::InvalidArgument(
+                    "dy and x must have the same shape: dy=",
+                    dy.shape().DebugString(), ", x=", x.shape().DebugString()));
+
+    const int axis = x.dims() - 1;
+    const int64 last_dim = x.dim_size(axis);
+
+    OP_REQUIRES(
+        ctx, gamma.NumElements() == last_dim,
+        errors::InvalidArgument("Gamma size mismatch: expected ", last_dim,
+                                ", got ", gamma.NumElements()));
+    OP_REQUIRES(
+        ctx, beta.NumElements() == last_dim,
+        errors::InvalidArgument("Beta size mismatch: expected ", last_dim,
+                                ", got ", beta.NumElements()));
+
+    Tensor* dx = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &dx));
+    Tensor* dgamma = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, gamma.shape(), &dgamma));
+    Tensor* dbeta = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, beta.shape(), &dbeta));
+
+    auto& handle = GetHandleByCtx(ctx);
+    auto stream = GetMusaStreamByCtx(ctx);
+
+    if (x.NumElements() == 0) {
+      musaMemsetAsync(const_cast<char*>(dx->tensor_data().data()), 0,
+                      dx->TotalBytes(), stream);
+      musaMemsetAsync(const_cast<char*>(dgamma->tensor_data().data()), 0,
+                      dgamma->TotalBytes(), stream);
+      musaMemsetAsync(const_cast<char*>(dbeta->tensor_data().data()), 0,
+                      dbeta->TotalBytes(), stream);
+      return;
+    }
+
+    TensorShape stat_shape;
+    for (int i = 0; i < x.dims() - 1; ++i) {
+      stat_shape.AddDim(x.dim_size(i));
+    }
+
+    Tensor y;
+    Tensor mean;
+    Tensor inv_var;
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(x.dtype(), x.shape(), &y));
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(x.dtype(), stat_shape, &mean));
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(x.dtype(), stat_shape, &inv_var));
+
+    mTensor mt_x = CreateMTensor(x, format_);
+    mTensor mt_gamma = CreateMTensor(gamma, format_);
+    mTensor mt_beta = CreateMTensor(beta, format_);
+    mTensor mt_y = CreateMTensor(y, format_);
+    mTensor mt_mean = CreateMTensor(mean, format_);
+    mTensor mt_inv_var = CreateMTensor(inv_var, format_);
+    mTensor mt_dy = CreateMTensor(dy, format_);
+    mTensor mt_dx = CreateMTensor(*dx, format_);
+    mTensor mt_dgamma = CreateMTensor(*dgamma, format_);
+    mTensor mt_dbeta = CreateMTensor(*dbeta, format_);
+
+    ::musa::dnn::LayerNorm ln;
+    ln.SetEpsilon(epsilon_);
+
+    std::vector<int> axis_vec;
+    axis_vec.push_back(axis);
+    ln.SetAxis(static_cast<int>(axis_vec.size()), axis_vec.data());
+
+    tensorflow::Allocator* tf_allocator =
+        ctx->device()->GetAllocator(tensorflow::AllocatorAttributes());
+    auto alloc_func =
+        [tf_allocator](
+            size_t size) -> std::unique_ptr<void, std::function<void(void*)>> {
+      void* ptr = tf_allocator->AllocateRaw(256, size);
+      return std::unique_ptr<void, std::function<void(void*)>>(
+          ptr, [tf_allocator](void* p) {
+            if (p) tf_allocator->DeallocateRaw(p);
+          });
+    };
+    ::musa::dnn::MemoryMaintainer maintainer(alloc_func);
+
+    mStatus status =
+        ln.Run(handle, mt_y, mt_mean, mt_inv_var, mt_x, mt_gamma, mt_beta,
+               maintainer);
+    OP_REQUIRES(
+        ctx, status == mStatus::SUCCESS,
+        errors::Internal("MUSA LayerNorm forward recompute failed. Status=",
+                         static_cast<int>(status)));
+
+    status = ln.RunBwd(handle, mt_dx, mt_dgamma, mt_dbeta, mt_dy, mt_x, mt_mean,
+                       mt_inv_var, mt_gamma, maintainer);
+    OP_REQUIRES(ctx, status == mStatus::SUCCESS,
+                errors::Internal("MUSA LayerNorm backward failed. Status=",
+                                 static_cast<int>(status)));
+  }
+
+ private:
+  float epsilon_;
+};
+
+#define REGISTER_MUSA_LAYERNORM_GRAD(TYPE)                                \
+  REGISTER_KERNEL_BUILDER(                                                 \
+      Name("MusaLayerNormGrad").Device("MUSA").TypeConstraint<TYPE>("T"), \
+      MusaLayerNormGradOp<TYPE>);
+
+REGISTER_MUSA_LAYERNORM_GRAD(float);
+REGISTER_MUSA_LAYERNORM_GRAD(Eigen::half);
+REGISTER_MUSA_LAYERNORM_GRAD(bfloat16);
+
+#undef REGISTER_MUSA_LAYERNORM_GRAD
+
+}  // namespace musa
+
+REGISTER_OP("MusaLayerNormGrad")
+    .Input("dy: T")
+    .Input("x: T")
+    .Input("gamma: T")
+    .Input("beta: T")
+    .Output("dx: T")
+    .Output("dgamma: T")
+    .Output("dbeta: T")
+    .Attr("T: {float, half, bfloat16}")
+    .Attr("epsilon: float = 0.00001")
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(1));
+      c->set_output(1, c->input(2));
+      c->set_output(2, c->input(3));
+      return OkStatus();
+    });
+
+}  // namespace tensorflow

python/__init__.py

@@ -59,7 +59,7 @@
     set_musa_allow_growth,
     set_musa_telemetry_config,
 )
-from . import ops, raw_ops
+from . import gradients, ops, raw_ops
 
 # Package version
 __version__ = "0.1.0"
@@ -80,6 +80,7 @@
 # Public API
 __all__ = [
     "__version__",
+    "gradients",
     "ops",
     "raw_ops",
     "load_plugin",

python/gradients.py

@@ -0,0 +1,51 @@
+# Copyright 2026 The TensorFlow MUSA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Gradient registrations for MUSA extension ops."""
+
+from tensorflow.python.framework import ops as tf_ops
+
+from . import raw_ops
+
+
+@tf_ops.RegisterGradient("MusaLayerNorm")
+def _musa_layer_norm_grad(op, grad):
+    epsilon = op.get_attr("epsilon")
+    dx, dgamma, dbeta = raw_ops.musa_layer_norm_grad(
+        dy=grad,
+        x=op.inputs[0],
+        gamma=op.inputs[1],
+        beta=op.inputs[2],
+        epsilon=epsilon,
+    )
+    return dx, dgamma, dbeta
+
+
+@tf_ops.RegisterGradient("MusaGelu")
+def _musa_gelu_grad(op, grad):
+    return raw_ops.musa_gelu_grad(
+        dy=grad,
+        x=op.inputs[0],
+        approximate=op.get_attr("approximate"),
+    )
+
+
+@tf_ops.RegisterGradient("MusaDropout")
+def _musa_dropout_grad(op, grad, _mask_grad):
+    return raw_ops.musa_dropout_grad(
+        grad=grad,
+        mask=op.outputs[1],
+        rate=op.get_attr("rate"),
+    )

python/ops.py

@@ -32,6 +32,17 @@ def layer_norm(x, gamma, beta, epsilon=0.00001, name=None):
     )
 
 
+def layer_norm_grad(dy, x, gamma, beta, epsilon=0.00001, name=None):
+    return raw_ops.musa_layer_norm_grad(
+        dy=dy,
+        x=x,
+        gamma=gamma,
+        beta=beta,
+        epsilon=epsilon,
+        name=name,
+    )
+
+
 def shifted_affine_map(data_left, mask, sliced_var_right, name=None):
     return raw_ops.musa_shifted_affine_map(
         data_left=data_left,
@@ -104,6 +115,15 @@ def gelu(x, approximate=False, name=None):
     )
 
 
+def gelu_grad(dy, x, approximate=False, name=None):
+    return raw_ops.musa_gelu_grad(
+        dy=dy,
+        x=x,
+        approximate=approximate,
+        name=name,
+    )
+
+
 def reshape_mat_mul(x, w, transpose_b=False, name=None):
     return raw_ops.musa_reshape_mat_mul(
         x=x,
@@ -158,9 +178,13 @@ def resource_apply_nadam(
     "clip",
     "dropout",
     "dropout_grad",
+    "dropout_grad",
     "gelu",
+    "gelu_grad",
     "interact",
     "layer_norm",
+    "layer_norm_grad",
+    "layer_norm",
     "matmul_bias_add",
     "reshape_mat_mul",
     "resource_apply_nadam",

test/musa_test_utils.py

@@ -16,6 +16,7 @@
 """Utilities for MUSA kernel tests."""
 
 import unittest
+
 import tensorflow as tf