destroy descriptor

Author: xgqdut2016

operatorspy/tests/random_sample.py

@@ -116,7 +116,7 @@ def test(lib, handle, torch_device, voc, random_val, topp, topk, temperature, x_
         torch.npu.synchronize()
     if torch_device == "sdaa":
         torch.sdaa.synchronize()
-    print(indices[0], data[indices[0]], ans, data[ans])
+    #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

@@ -79,7 +79,10 @@ void topkKernel(const void *probs, void *index, void *value, int topk, int voc,
     status = tecodnnTopkEx(handle, axis, topk, largest, sorted, input_desc_teco, probs, value_desc_teco, value,
                            index_desc_teco, index, compute_workspace, workSpaceSizeInBytes);
     sdaaStreamSynchronize(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));
     }
@@ -106,7 +109,8 @@ void softmaxKernel(const void *probs, void *destination, int voc, tecodnnHandle_
 
     status = tecodnnSoftmaxForward(handle, algo, mode, &alpha, x_desc_teco, probs, &beta, y_desc_teco, destination);
     sdaaStreamSynchronize(stream);
-
+    tecodnnDestroyTensorDescriptor(x_desc_teco);
+    tecodnnDestroyTensorDescriptor(y_desc_teco);
     if (status != TECODNN_STATUS_SUCCESS) {
         printf("softmax %s\n", tecodnnGetErrorString(status));
     }
@@ -145,7 +149,8 @@ void cumSumKernel(void *value, void *scan_value, int topk_, tecodnnHandle_t hand
 
     status = tecodnnCumSum(handle, 3, a_desc_teco, value, c_desc_teco, scan_value);
     sdaaStreamSynchronize(stream);
-
+    tecodnnDestroyTensorDescriptor(a_desc_teco);
+    tecodnnDestroyTensorDescriptor(c_desc_teco);
     if (status != TECODNN_STATUS_SUCCESS) {
         printf("scan %s\n", tecodnnGetErrorString(status));
     }
@@ -176,9 +181,14 @@ __global__ void sample(T *scan_value, int64_t *index, uint64_t *result, float ra
             }
         }
     }
-    
+
 }
 
+__global__ void randomSampleKernel(uint64_t *result, int64_t *index){
+    if(threadIdx == 0){
+        result[0] = index[0];
+    }
+}
 infiniopStatus_t tecoRandomSample(RandomSampleTecoDescriptor_t desc,
                                   void *workspace,
                                   uint64_t workspace_size,
@@ -204,26 +214,34 @@ 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);
-        softmaxKernel<half>(probs, (void *) destination, desc->voc, desc->handle, desc->stream);
-
         sdaaEvent_t event; 
         sdaaEventCreate(&event); 
+        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();
+            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);
+            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();
+            sdaaEventRecord(event, desc->stream);
+            sample<half><<<1, desc->stream>>>(scan_value, index, (uint64_t *) result, random_val, topp, topk);
+            
+            sdaaEventSynchronize(event);//必须使用这些阻塞，确保kernel完成
+            sdaaDeviceSynchronize();
+
+            
+        }
+        else {
+            sdaaEventRecord(event, desc->stream);
+            randomSampleKernel<<<1, desc->stream>>>((uint64_t *)result, index);
+            sdaaEventSynchronize(event);
+            sdaaDeviceSynchronize();
 
+        }
         sdaaEventDestroy(event);
-        
 
         return STATUS_SUCCESS;
     }