random sample算子太初平台

operatorspy/tests/random_sample.py

@@ -31,47 +31,32 @@ class RandomSampleDescriptor(Structure):
 
 def random_sample(data, random_val, topp, topk, voc, temperature, torch_device):
     indices = torch.zeros([topk], dtype = torch.int64)
-    dataNp = data.clone().detach()
-    sorted_indices = torch.arange(voc)
-
-    for i in range(topk):
-        for j in range(i + 1, voc):
-            if(dataNp[i] < dataNp[j]):
-                tmp = dataNp[i].clone().detach()
-                dataNp[i] = dataNp[j].clone().detach()
-                dataNp[j] = tmp
-
-                tmpInd = sorted_indices[i].clone().detach()
-                sorted_indices[i] = sorted_indices[j].clone().detach()
-                sorted_indices[j] = tmpInd
+    dataNp = data.clone()
 
-    #sorted_indices = torch.argsort(dataNp, descending=True)
+    sorted_indices = torch.argsort(dataNp, descending=True)
     indices = sorted_indices[:topk] 
 
     dataNp = dataNp[sorted_indices]
 
     globalM = dataNp[0]
     dataNp = (dataNp - globalM) / temperature
     dataNp = torch.softmax(dataNp.float(), dim = 0)
-    sum_s = 0
+
+    for i in range(1, topk):
+        dataNp[i] = dataNp[i] + dataNp[i - 1]
+
     for end in range(topk):
-        sum_s += dataNp[end]
-        if(sum_s >= topp):
+        if(dataNp[end] >= topp):
             break
     if(end < topk - 1):
         end += 1
     else:
         end = topk
 
-    sum_s = 0
-    for i in range(end):
-        sum_s += dataNp[i]
-    random_val *= sum_s
+    random_val *= dataNp[end - 1]
 
-    sum_s = 0
     for i in range(end):
-        sum_s += dataNp[i]
-        if(random_val < sum_s):
+        if(random_val < dataNp[i]):
             return indices[i]
 
 def random_sample_0(data):
@@ -129,7 +114,7 @@ def test(lib, handle, torch_device, voc, random_val, topp, topk, temperature, x_
     )
     if torch_device == "npu":
         torch.npu.synchronize()
-
+    
     assert indices[0].type(ans.dtype) == ans or data[ans] == data[indices[0]]
     check_error(lib.infiniopDestroyRandomSampleDescriptor(descriptor))
     print("Test passed!")
@@ -168,7 +153,13 @@ def test_ascend(lib, test_cases):
         test(lib, handle, "npu", voc, random_val, topp, topk, temperature)
     destroy_handle(lib, handle) 
 
-
+def test_teco(lib, test_cases):
+    import torch_sdaa
+    device = DeviceEnum.DEVICE_TECO
+    handle = create_handle(lib, device)
+    for (voc, random_val, topp, topk, temperature) in test_cases:
+        test(lib, handle, "sdaa", voc, random_val, topp, topk, temperature)
+    destroy_handle(lib, handle)
 
 if __name__ == "__main__":
     test_cases = [
@@ -224,6 +215,9 @@ def test_ascend(lib, test_cases):
         test_bang(lib, test_cases)
     if args.ascend:
         test_ascend(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend):
+    if args.teco:
+        test_teco(lib, test_cases)
+
+    if not (args.cpu or args.cuda or args.bang or args.ascend or args.teco):
         test_cpu(lib, test_cases)
     print("\033[92mTest passed!\033[0m")

src/ops/random_sample/operator.cc

@@ -14,6 +14,9 @@
 #ifdef ENABLE_ASCEND_NPU
 #include "ascend/random_sample.h"
 #endif
+#ifdef ENABLE_TECO_SDAA
+#include "teco/random_sample_teco.h"
+#endif
 
 __C infiniopStatus_t infiniopCreateRandomSampleDescriptor(infiniopHandle_t handle, infiniopRandomSampleDescriptor_t *desc_ptr, infiniopTensorDescriptor_t result, infiniopTensorDescriptor_t probs) {
     switch (handle->device) {
@@ -35,8 +38,14 @@ __C infiniopStatus_t infiniopCreateRandomSampleDescriptor(infiniopHandle_t handl
 #ifdef ENABLE_ASCEND_NPU
         case DevAscendNpu: {
             return ascendCreateRandomSampleDescriptor((AscendHandle_t) handle,
-                                                     (RandomSampleAscendDescriptor_t *) desc_ptr, result, probs);
+                                                      (RandomSampleAscendDescriptor_t *) desc_ptr, result, probs);
         }
+#endif
+#ifdef ENABLE_TECO_SDAA
+        case DevTecoSDAA:
+            return tecoCreateRandomSampleDescriptor((TecoHandle_t) handle,
+                                                    (RandomSampleTecoDescriptor_t *) desc_ptr, result, probs);
+            ;
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -64,6 +73,10 @@ __C infiniopStatus_t infiniopGetRandomSampleWorkspaceSize(infiniopRandomSampleDe
         case DevAscendNpu: {
             return ascendGetRandomSampleWorkspaceSize((RandomSampleAscendDescriptor_t) desc, size);
         }
+#endif
+#ifdef ENABLE_TECO_SDAA
+        case DevTecoSDAA:
+            return tecoGetRandomSampleWorkspaceSize((RandomSampleTecoDescriptor_t) desc, size);
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -97,6 +110,10 @@ __C infiniopStatus_t infiniopRandomSample(infiniopRandomSampleDescriptor_t desc,
         case DevAscendNpu: {
             return ascendRandomSample((RandomSampleAscendDescriptor_t) desc, workspace, workspace_size, result, probs, random_val, topp, topk, temperature, stream);
         }
+#endif
+#ifdef ENABLE_TECO_SDAA
+        case DevTecoSDAA:
+            return tecoRandomSample((RandomSampleTecoDescriptor_t) desc, workspace, workspace_size, result, probs, random_val, topp, topk, temperature, stream);
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -121,6 +138,10 @@ __C infiniopStatus_t infiniopDestroyRandomSampleDescriptor(infiniopRandomSampleD
         case DevAscendNpu: {
             return ascendDestroyRandomSampleDescriptor((RandomSampleAscendDescriptor_t) desc);
         }
+#endif
+#ifdef ENABLE_TECO_SDAA
+        case DevTecoSDAA:
+            return tecoDestroyRandomSampleDescriptor((RandomSampleTecoDescriptor_t) desc);
 #endif
     }
     return STATUS_BAD_DEVICE;

src/ops/random_sample/teco/random_sample_teco.h

@@ -0,0 +1,42 @@
+#ifndef __SDAA_RANDOM_SAMPLE_H__
+#define __SDAA_RANDOM_SAMPLE_H__
+
+#include "../../../devices/teco/teco_handle.h"
+#include "../../utils.h"
+#include "operators.h"
+#include <sdaa_runtime.h>
+
+struct RandomSampleTecoDescriptor {
+    Device device;
+    int device_id;
+    tecodnnHandle_t handle;
+    sdaaStream_t stream;
+    DT dtype;
+    int voc;
+    DT rDtype;
+    int rLength;
+};
+
+typedef struct RandomSampleTecoDescriptor *RandomSampleTecoDescriptor_t;
+
+infiniopStatus_t tecoCreateRandomSampleDescriptor(TecoHandle_t handle,
+                                                  RandomSampleTecoDescriptor_t *desc_ptr, infiniopTensorDescriptor_t result,
+                                                  infiniopTensorDescriptor_t probs);
+
+infiniopStatus_t tecoGetRandomSampleWorkspaceSize(RandomSampleTecoDescriptor_t desc, uint64_t *size);
+
+infiniopStatus_t tecoRandomSample(RandomSampleTecoDescriptor_t desc,
+                                  void *workspace,
+                                  uint64_t workspace_size,
+                                  void *result,
+                                  void const *probs,
+                                  float random_val,
+                                  float topp,
+                                  int topk,
+                                  float temperature,
+                                  void *stream);
+
+infiniopStatus_t tecoDestroyRandomSampleDescriptor(RandomSampleTecoDescriptor_t desc);
+
+
+#endif