From 23486328026ef98374d6f6cc59d0ae13cd47f487 Mon Sep 17 00:00:00 2001
From: LindseyMei <648816901@qq.com>
Date: Thu, 25 Jun 2026 06:39:29 +0000
Subject: [PATCH] feat: support sigmoid operator on metax

Add MetaX backend for the sigmoid elementwise operator, reusing the existing cuda::SigmoidOp kernel through the elementwise MetaX descriptor (mirrors the silu metax implementation). Supports BF16/F16/F32/F64.

Verified with test/infiniop/sigmoid.py --metax on MetaX C500: passes accuracy check against torch.sigmoid across shapes/strides and inplace/out-of-place for F16/F32/BF16.

Signed-off-by: LindseyMei <648816901@qq.com>
---
 .../ops/sigmoid/metax/sigmoid_metax.h         |  8 +++
 .../ops/sigmoid/metax/sigmoid_metax.maca      | 60 +++++++++++++++++++
 src/infiniop/ops/sigmoid/operator.cc          | 15 +++++
 3 files changed, 83 insertions(+)
 create mode 100644 src/infiniop/ops/sigmoid/metax/sigmoid_metax.h
 create mode 100644 src/infiniop/ops/sigmoid/metax/sigmoid_metax.maca

diff --git a/src/infiniop/ops/sigmoid/metax/sigmoid_metax.h b/src/infiniop/ops/sigmoid/metax/sigmoid_metax.h
new file mode 100644
index 000000000..76189a759
--- /dev/null
+++ b/src/infiniop/ops/sigmoid/metax/sigmoid_metax.h
@@ -0,0 +1,8 @@
+#ifndef __SIGMOID_METAX_API_H__
+#define __SIGMOID_METAX_API_H__
+
+#include "../../../elementwise/metax/elementwise_metax_api.h"
+
+ELEMENTWISE_DESCRIPTOR(sigmoid, metax)
+
+#endif // __SIGMOID_METAX_API_H__
diff --git a/src/infiniop/ops/sigmoid/metax/sigmoid_metax.maca b/src/infiniop/ops/sigmoid/metax/sigmoid_metax.maca
new file mode 100644
index 000000000..dcda5894c
--- /dev/null
+++ b/src/infiniop/ops/sigmoid/metax/sigmoid_metax.maca
@@ -0,0 +1,60 @@
+#include "sigmoid_metax.h"
+
+#include "../../../elementwise/metax/elementwise_metax.h"
+
+#include "../cuda/kernel.cuh"
+
+namespace op::sigmoid::metax {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    std::vector<infiniopTensorDescriptor_t> input_desc_vec) {
+
+    auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
+    auto dtype = out_desc->dtype();
+
+    const auto &input_desc = input_desc_vec.at(0);
+    const auto &output_shape = out_desc->shape();
+    const auto &input_shape = input_desc->shape();
+
+    CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64);
+
+    CHECK_SAME_SHAPE(output_shape, input_shape);
+
+    // create METAX elementwise descriptor
+    CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec)
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    std::vector<const void *> inputs,
+    void *stream) const {
+
+    if (workspace_size < _workspace_size) {
+        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
+    }
+
+    switch (_dtype) {
+    case INFINI_DTYPE_BF16:
+        return _device_info->calculate<256, cuda::SigmoidOp, cuda_bfloat16>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_F16:
+        return _device_info->calculate<256, cuda::SigmoidOp, half>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_F32:
+        return _device_info->calculate<256, cuda::SigmoidOp, float>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_F64:
+        return _device_info->calculate<256, cuda::SigmoidOp, double>(_info, workspace, output, inputs, stream);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace op::sigmoid::metax
diff --git a/src/infiniop/ops/sigmoid/operator.cc b/src/infiniop/ops/sigmoid/operator.cc
index 6d416b130..89146d820 100644
--- a/src/infiniop/ops/sigmoid/operator.cc
+++ b/src/infiniop/ops/sigmoid/operator.cc
@@ -8,6 +8,9 @@
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
 #include "nvidia/sigmoid_nvidia.cuh"
 #endif
+#ifdef ENABLE_METAX_API
+#include "metax/sigmoid_metax.h"
+#endif
 
 __INFINI_C infiniStatus_t infiniopCreateSigmoidDescriptor(
     infiniopHandle_t handle,
@@ -43,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreateSigmoidDescriptor(
 #ifdef ENABLE_HYGON_API
         CREATE(INFINI_DEVICE_HYGON, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        CREATE(INFINI_DEVICE_METAX, metax);
+#endif
 
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -76,6 +82,9 @@ __INFINI_C infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescrip
 #endif
 #ifdef ENABLE_HYGON_API
         GET(INFINI_DEVICE_HYGON, nvidia)
+#endif
+#ifdef ENABLE_METAX_API
+        GET(INFINI_DEVICE_METAX, metax)
 #endif
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -118,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopSigmoid(
 #ifdef ENABLE_HYGON_API
         CALCULATE(INFINI_DEVICE_HYGON, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        CALCULATE(INFINI_DEVICE_METAX, metax);
+#endif
 
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -154,6 +166,9 @@ infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc) {
 #ifdef ENABLE_HYGON_API
         DELETE(INFINI_DEVICE_HYGON, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        DELETE(INFINI_DEVICE_METAX, metax);
+#endif
 
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;