implement swiglu inplace and out of place operation

Author: crapromer

operatorspy/tests/swiglu.py

@@ -124,7 +124,6 @@ def test_in_place1(
             descriptor, a_tensor.data, a_tensor.data, b_tensor.data, None
         )
     )
-
     assert torch.allclose(a, ans, atol=1e-4, rtol=1e-2)
     print("in-place1 Test passed!")
 

src/devices/teco/tensor_teco.cpp

@@ -15,7 +15,7 @@
 #include "ascend/swiglu.h"
 #endif
 #ifdef ENABLE_TECO_SDAA
-#include "teco/swiglu_tecodnn.h"
+#include "teco/swiglu_sdaa.h"
 #endif
 
 __C infiniopStatus_t infiniopCreateSwiGLUDescriptor(infiniopHandle_t handle,

src/devices/teco/tensor_teco.h

@@ -1,24 +1,22 @@
-#ifndef __TECO_SWIGLU_H__
-#define __TECO_SWIGLU_H__
-
+#ifndef __SDAA_SWIGLU_H__
+#define __SDAA_SWIGLU_H__
 #include "operators.h"
 #include <sdaa_runtime.h>
-#include <tecodnn.h>
 #include "../../../devices/teco/teco_handle.h"
 struct SwiGLUTecoDescriptor {
     Device device;
     int device_id;
     sdaaStream_t stream;
-    tecodnnHandle_t handle;
-    tecodnnActivationDescriptor_t activationDesc;
-    tecodnnTensorDescriptor_t aDesc,bDesc,cDesc;
+    uint64_t rows,cols;
+    int64_t lda,ldb,ldc;
 };
 
 typedef struct SwiGLUTecoDescriptor *SwiGLUTecoDescriptor_t;
 
+
 infiniopStatus_t tecoCreateSwiGLUDescriptor(TecoHandle_t handle,
                                              SwiGLUTecoDescriptor_t *desc_ptr,
-                                             infiniopTensorDescriptor_t c_desc,                                             
+                                             infiniopTensorDescriptor_t c_desc,
                                              infiniopTensorDescriptor_t a_desc,
                                              infiniopTensorDescriptor_t b_desc);
 
@@ -30,4 +28,5 @@ infiniopStatus_t tecoSwiGLU(SwiGLUTecoDescriptor_t desc,
 
 infiniopStatus_t tecoDestroySwiGLUDescriptor(SwiGLUTecoDescriptor_t desc);
 
+
 #endif

src/ops/swiglu/operator.cc

@@ -0,0 +1,70 @@
+#include "swiglu_sdaa.h"
+__local__ halfv16 tempa, tempb, tempc;
+
+__device__ void silu_halfv16(halfv16 *c, halfv16 *a, halfv16 *b) {
+    floatv16 one_v = simd_stretch(1.0f);
+    floatv16 a_silu = simd_div(simd_cvt_h2f(*b), simd_add(one_v, simd_exp(0 - simd_cvt_h2f(*b))));
+    halfv16 out = simd_cvt_f2h(simd_mul(simd_cvt_h2f(*a), a_silu));
+    *c = out;
+}
+
+__device__ void silu_half(half *c, const half *a, const half *b) {
+    *c = (*b) * (*a)/ (1.0 + expf(0 - *b)) ;
+}
+
+__global__ void swiglu(half *c, half const *a, half const *b, size_t rows, size_t cols, size_t lda, size_t ldb, size_t ldc) {
+    int vector_size = 16;
+    for (size_t i = 0; i < rows / threadDim + 1; i++) {
+        if (threadIdx < rows - i * threadDim) {
+            size_t j = 0;
+            for (; j < cols / vector_size; j++) {
+                simd_load(tempa, a + (threadIdx + i * threadDim) * lda + j * vector_size);
+                simd_load(tempb, b + (threadIdx + i * threadDim) * ldb + j * vector_size);
+                silu_halfv16(&tempc, &tempa, &tempb);
+                simd_store(tempc, c + (threadIdx + i * threadDim) * ldc + j * vector_size);
+            }
+            for (size_t k = 0; k < cols - j * vector_size; k++)
+            {
+                silu_half(
+                c + (threadIdx + i * threadDim) * ldc + j * vector_size + k,
+                a + (threadIdx + i * threadDim) * lda + j * vector_size + k,
+                b + (threadIdx + i * threadDim) * ldb + j * vector_size + k);
+            }
+            
+        }
+    }
+}
+
+infiniopStatus_t tecoCreateSwiGLUDescriptor(TecoHandle_t handle,
+                                            SwiGLUTecoDescriptor_t *desc_ptr,
+                                            infiniopTensorDescriptor_t c_desc,
+                                            infiniopTensorDescriptor_t a_desc,
+                                            infiniopTensorDescriptor_t b_desc) {
+    *desc_ptr = new SwiGLUTecoDescriptor{
+        handle->device,
+        handle->device_id,
+        handle->stream,
+        a_desc->shape[0],
+        a_desc->shape[1],
+        a_desc->strides[0],
+        b_desc->strides[0],
+        c_desc->strides[0],
+    };
+    return STATUS_SUCCESS;
+}
+
+infiniopStatus_t tecoSwiGLU(SwiGLUTecoDescriptor_t desc,
+                            void *c,
+                            void const *a,
+                            void const *b,
+                            void *stream) {
+    auto a_ptr = reinterpret_cast<const half *>(a);
+    auto b_ptr = reinterpret_cast<const half *>(b);
+    auto c_ptr = reinterpret_cast<half *>(c);
+    swiglu<<<1>>>(c_ptr, a_ptr, b_ptr, desc->rows, desc->cols, desc->lda, desc->ldb, desc->ldc);
+    return STATUS_SUCCESS;
+}
+
+infiniopStatus_t tecoDestroySwiGLUDescriptor(SwiGLUTecoDescriptor_t desc) {
+    return STATUS_SUCCESS;
+}

src/ops/swiglu/teco/swiglu_tecodnn.h renamed to src/ops/swiglu/teco/swiglu_sdaa.h

@@ -243,7 +243,7 @@ if has_config("teco") then
             local cc = "/opt/tecoai/bin/tecocc"
 
             local includedirs = table.concat(target:get("includedirs"), " ")
-            local args = {sourcefile, "-o", objectfile}
+            local args = {sourcefile, "-o", objectfile,"-O2", "-fPIC", "-Wall", "-Werror", "-std=c++17", "-pthread","-c"}
 
             for _, includedir in ipairs(target:get("includedirs")) do
                 table.insert(args, "-I" .. includedir)