[Ascend] fuj/fix-generator-with-getSeedAndOffset (#1233)

fix generator with getSeedAndOffset
DeepLink-org · Jun 11, 2024 · 532bb11 · 532bb11
1 parent ec68771
commit 532bb11
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 17 deletions.
diff --git a/impl/ascend/functions/dropout.cpp b/impl/ascend/functions/dropout.cpp
@@ -23,8 +23,9 @@ diopiError_t npuDropoutOut(diopiContextHandle_t ctx, diopiTensorHandle_t out, di
     diopiError_t ret = diopiRequireTensor(ctx, &maskNpu, &maskSize, nullptr, diopi_dtype_uint8, diopi_device);
     ASCEND_CHECK_ABORT(ret == diopiSuccess, "[npuDropoutOut] require tensor for mask failed.");
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     DIOPI_ASCEND_CALL_ACLNN(aclnnDropoutGenMask, ctx, inAt.shape(), p, seed, offset, maskNpu);
     DIOPI_ASCEND_CALL_ACLNN(aclnnDropoutDoMask, ctx, input, maskNpu, p, out);
@@ -57,8 +58,9 @@ diopiError_t npuDropout2dOut(diopiContextHandle_t ctx, diopiTensorHandle_t out,
     diopiError_t ret = diopiRequireTensor(ctx, &maskNpu, &maskNpuSize, nullptr, diopi_dtype_uint8, diopi_device);
     ASCEND_CHECK_ABORT(ret == diopiSuccess, "[npuDropout2dOut] require tensor for mask failed.");
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     DIOPI_ASCEND_CALL_ACLNN(aclnnDropoutGenMask, ctx, inAt.shape(), p, seed, offset, maskNpu);
     DIOPI_ASCEND_CALL_ACLNN(aclnnDropoutDoMask, ctx, input2d, maskNpu, p, out2d);

diff --git a/impl/ascend/functions/normal.cpp b/impl/ascend/functions/normal.cpp
@@ -16,8 +16,9 @@ diopiError_t diopiNormal(diopiContextHandle_t ctx, diopiTensorHandle_t out, doub
         return diopiSuccess;
     }
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     float meanCast = static_cast<float>(mean);
     float rstdCast = static_cast<float>(std);
@@ -26,8 +27,9 @@ diopiError_t diopiNormal(diopiContextHandle_t ctx, diopiTensorHandle_t out, doub
 }
 
 diopiError_t diopiNormalInp(diopiContextHandle_t ctx, diopiTensorHandle_t inout, double mean, double std, diopiGeneratorHandle_t generator) {
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     float meanCast = static_cast<float>(mean);
     float rstdCast = static_cast<float>(std);
@@ -42,8 +44,9 @@ diopiError_t diopiNormalTensor(diopiContextHandle_t ctx, diopiTensorHandle_t out
         return diopiSuccess;
     }
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     DIOPI_ASCEND_CALL_ACLNN(aclnnNormalTensorTensor, ctx, mean, std, seed, offset, out);
     return diopiSuccess;
@@ -56,8 +59,9 @@ diopiError_t diopiNormalScalarTensor(diopiContextHandle_t ctx, diopiTensorHandle
         return diopiSuccess;
     }
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     float meanCast = static_cast<float>(mean);
     DIOPI_ASCEND_CALL_ACLNN(aclnnNormalFloatTensor, ctx, meanCast, std, seed, offset, out);
@@ -71,8 +75,9 @@ diopiError_t diopiNormalTensorScalar(diopiContextHandle_t ctx, diopiTensorHandle
         return diopiSuccess;
     }
 
-    uint64_t seed, offset;
-    DIOPI_CALL(diopiGeneratorGetSeedAndOffset(generator, &seed, &offset));
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
 
     float rstdCast = static_cast<float>(std);
     DIOPI_ASCEND_CALL_ACLNN(aclnnNormalTensorFloat, ctx, mean, rstdCast, seed, offset, out);

diff --git a/impl/ascend/functions/uniform.cpp b/impl/ascend/functions/uniform.cpp
@@ -11,9 +11,9 @@ namespace impl {
 namespace ascend {
 
 diopiError_t diopiUniformInp(diopiContextHandle_t ctx, diopiTensorHandle_t inout, double from, double to, diopiGeneratorHandle_t generator) {
-    uint64_t seed = 0;
-    uint64_t offset = 0;
-    diopiGeneratorGetSeedAndOffset(generator, &seed, &offset);
+    const std::pair<uint64_t, uint64_t> gen = getSeedAndOffset(ctx, generator, 10);
+    const uint64_t seed = gen.first;
+    const uint64_t offset = gen.second;
     DIOPI_ASCEND_CALL_ACLNN(aclnnInplaceUniform, ctx, inout, from, to, seed, offset);
     return diopiSuccess;
 }