#include <torch/extension.h>
#include <torch/library.h>
#include <torch/version.h>
#include <torch_npu/csrc/core/npu/NPUStream.h>
#include <torch_npu/csrc/framework/OpCommand.h>
#include <torch_npu/csrc/npu/Module.h>
#include "utils.h"
/*
 * How to write a meta implementation for a custom operator (meta kernel):
 *
 * Meta implementations are used for shape and dtype inference, tracing, and export.
 * They do NOT perform any real computation or allocate device memory.
 * Instead, they return empty tensors with the correct shapes, dtypes, and device types.
 *
 * Steps to write a meta implementation:
 * 1. The function signature should match the operator's schema, but only use the arguments
 *    necessary to infer output shapes and dtypes.
 * 2. Use input tensor shapes, dtypes, and any relevant arguments to compute the output shapes.
 * 3. Return empty tensors (e.g., at::empty_symint, at::empty_like) with the correct shape and dtype.
 * 4. Do NOT perform any real computation or data movement.
 * 5. Register the meta implementation with the "Meta" dispatch key using TORCH_LIBRARY_IMPL or similar.
 *
 * Example:
 *   std::tuple<at::Tensor, at::Tensor> my_op_meta(
 *       at::Tensor &input, int64_t some_param) {
 *     // Infer output shape based on input and parameters
 *     auto out_shape = ...;
 *     at::Tensor out = at::empty_symint(out_shape, input.options());
 *     // Return empty tensor(s) with correct shape/dtype
 *     return {out, ...};
 *   }
 *
 * See below for real examples.
 */

namespace mindspeed_ops {
namespace meta {

std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_add_rms_norm_bias_meta(const at::Tensor &x1, const at::Tensor &x2,
    const at::Tensor &gamma, const c10::optional<at::Tensor> &beta, double epsilon) {
    int64_t dim_x = x1.dim();
    int64_t dim_gamma = gamma.dim();
    int64_t diff = dim_x - dim_gamma;
    c10::SymDimVector new_shape;
    at::Tensor rstd;

    if (diff > 0) {
        new_shape.reserve(dim_x);
        auto x1_sizes = x1.sym_sizes();
        for (int64_t i = 0; i < diff; ++i) {
            new_shape.push_back(x1_sizes[i]);
        }
        for (int64_t i = 0; i < dim_gamma; ++i) {
            new_shape.push_back(c10::SymInt(1));
        }
    } else {
        new_shape.assign(dim_x, c10::SymInt(1));
    }
    rstd = at::empty_symint(new_shape, x1.options().dtype(at::kFloat));
    at::Tensor y = at::empty_symint(x1.sym_sizes(), x1.options());
    at::Tensor x = at::empty_symint(x1.sym_sizes(), x1.options());
    return std::tuple<at::Tensor, at::Tensor, at::Tensor>(y, rstd, x);
}
} // namespace meta
} // namespace mindspeed_ops

// Register the meta implementations of the custom kernels for symbolic tracing, this will also
// Pybind on platform
namespace {
TORCH_LIBRARY_IMPL_EXPAND(CONCAT(_C, _ascend), Meta, ops) {
    // Add_Rms_Norm_Bias
    ops.impl("npu_add_rms_norm_bias", &mindspeed_ops::meta::npu_add_rms_norm_bias_meta);
}
} // namespace
#endif