success random sample

Author: xgqdut2016

operatorspy/tests/random_sample.py

@@ -114,9 +114,7 @@ def test(lib, handle, torch_device, voc, random_val, topp, topk, temperature, x_
     )
     if torch_device == "npu":
         torch.npu.synchronize()
-    if torch_device == "sdaa":
-        torch.sdaa.synchronize()
-    #print(indices[0], data[indices[0]], ans, data[ans])
+    
     assert indices[0].type(ans.dtype) == ans or data[ans] == data[indices[0]]
     check_error(lib.infiniopDestroyRandomSampleDescriptor(descriptor))
     print("Test passed!")

src/ops/random_sample/teco/random_sample_teco.scpp

@@ -41,6 +41,7 @@ infiniopStatus_t tecoGetRandomSampleWorkspaceSize(RandomSampleTecoDescriptor_t d
 }
 
 infiniopStatus_t tecoDestroyRandomSampleDescriptor(RandomSampleTecoDescriptor_t desc) {
+    //sdaaStreamDestroy(desc->stream);
     delete desc;
     return STATUS_SUCCESS;
 }
@@ -213,35 +214,58 @@ infiniopStatus_t tecoRandomSample(RandomSampleTecoDescriptor_t desc,
 
 
         tecodnnMemset(desc->handle, value, 0, topk_);
-        topkKernel<half>(probs, (void *)index, (void *)value, topk, desc->voc, desc->handle, desc->stream);
-        sdaaEvent_t event; 
-        sdaaEventCreate(&event); 
+        
+        int voc = desc->voc;
+        tecodnnSetStream(desc->handle, desc->stream);
+        tecodnnStatus_t status;
+
+        tecodnnTensorDescriptor_t input_desc_teco, value_desc_teco, index_desc_teco;
+        tecodnnCreateTensorDescriptor(&input_desc_teco);
+        tecodnnCreateTensorDescriptor(&value_desc_teco);
+        tecodnnCreateTensorDescriptor(&index_desc_teco);
+
+        int32_t probsDims[2] = {1, voc}, probsStrides[2] = {voc, 1};
+        int32_t resultDims[2] = {1, topk}, resultStrides[2] = {topk, 1};
+        tecodnnSetTensorNdDescriptor(input_desc_teco, TECODNN_DATA_HALF, 2, probsDims, probsStrides);
+        tecodnnSetTensorNdDescriptor(value_desc_teco, TECODNN_DATA_HALF, 2, resultDims, resultStrides);
+
+        tecodnnSetTensorNdDescriptor(index_desc_teco, TECODNN_DATA_INT64, 2, resultDims, resultStrides);
+
+        size_t workSpaceSizeInBytes;
+        int axis = 1;
+        bool largest = true;
+        bool sorted = true;
+        tecodnnGetTopkExWorkspaceSize(desc->handle, axis, topk, largest, sorted, input_desc_teco, value_desc_teco,
+                                    index_desc_teco, &workSpaceSizeInBytes);
+        void *compute_workspace;
+        sdaaMalloc((void **) &compute_workspace, workSpaceSizeInBytes);
+
+        status = tecodnnTopkEx(desc->handle, axis, topk, largest, sorted, input_desc_teco, probs, value_desc_teco, value,
+                            index_desc_teco, index, compute_workspace, workSpaceSizeInBytes);
+        sdaaStreamSynchronize(desc->stream);
+        sdaaFree(compute_workspace);
+        tecodnnDestroyTensorDescriptor(input_desc_teco);
+        tecodnnDestroyTensorDescriptor(value_desc_teco);
+        tecodnnDestroyTensorDescriptor(index_desc_teco);
+        if (status != TECODNN_STATUS_SUCCESS) {
+            printf("topk %s\n", tecodnnGetErrorString(status));
+        }
+        
+        //topkKernel<half>(probs, (void *)index, (void *)value, topk, desc->voc, desc->handle, desc->stream);
+        
         if (topp > 0 && topk > 1){
             softmaxKernel<half>(probs, (void *) destination, desc->voc, desc->handle, desc->stream);
-
-            sdaaEventRecord(event, desc->stream);
             memKernel<half><<<1, desc->stream>>>(destination, value, index, topk);
-            sdaaEventSynchronize(event);
             sdaaDeviceSynchronize();
-
             cumSumKernel<half>((void *) value, (void *) scan_value, topk_, desc->handle, desc->stream);
-
-            sdaaEventRecord(event, desc->stream);
             sample<half><<<1, desc->stream>>>(scan_value, index, (uint64_t *) result, random_val, topp, topk);
-            
-            sdaaEventSynchronize(event);//必须使用这些阻塞，确保kernel完成
-            sdaaDeviceSynchronize();
-
-            
+            sdaaDeviceSynchronize(); 
         }
         else {
-            sdaaEventRecord(event, desc->stream);
             randomSampleKernel<<<1, desc->stream>>>((uint64_t *)result, index);
-            sdaaEventSynchronize(event);
             sdaaDeviceSynchronize();
-
         }
-        sdaaEventDestroy(event);
+        
 
         return STATUS_SUCCESS;
     }