From d45915a8a03af12d4348a06bd288d326ac28489a Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Thu, 4 Jan 2024 19:24:36 +0200 Subject: [PATCH] BINDINGS/GO: Locked OS thread before setting CUDA device. --- bindings/go/src/examples/perftest/perftest.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bindings/go/src/examples/perftest/perftest.go b/bindings/go/src/examples/perftest/perftest.go index df72fa66431..468f075e593 100644 --- a/bindings/go/src/examples/perftest/perftest.go +++ b/bindings/go/src/examples/perftest/perftest.go @@ -17,6 +17,7 @@ import ( . "ucx" "unsafe" . "cuda" + "runtime" ) type PerfTestParams struct { @@ -92,6 +93,7 @@ func initContext() { func tryCudaSetDevice() { if perfTestParams.memType == UCS_MEMORY_TYPE_CUDA { + runtime.LockOSThread() if ret := CudaSetDevice(); ret != nil { panic(ret) } @@ -233,8 +235,6 @@ func serverAmRecvHandler(header unsafe.Pointer, headerSize uint64, data *UcpAmDa if data.IsDataValid() { atomic.AddUint32(&perfTest.numCompletedRequests, 1) } else { - tryCudaSetDevice() - data.Receive(getAddressOffsetForThread(tid), perfTestParams.messageSize, (&UcpRequestParams{}).SetMemType(perfTestParams.memType).SetCallback(func(request *UcpRequest, status UcsStatus, length uint64) { atomic.AddUint32(&perfTest.numCompletedRequests, 1) @@ -281,12 +281,12 @@ func serverStart() error { } func clientThreadDoIter(i int, t uint) { + tryCudaSetDevice() + start := time.Now() var request *UcpRequest requestParams := (&UcpRequestParams{}).SetMemType(perfTestParams.memType) - tryCudaSetDevice() - header := unsafe.Pointer(&t) request, _ = perfTest.eps[t].SendAmNonBlocking(t, header, uint64(unsafe.Sizeof(t)), getAddressOffsetForThread(t), perfTestParams.messageSize, 0, requestParams)