1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
/*
* Copyright (C) 2021, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
* See file LICENSE for terms.
*/
package goucxtests
import (
. "github.com/openucx/ucx/bindings/go/src/ucx"
"unsafe"
)
const selfEpTag uint64 = ^uint64(0)
func memoryAllocate(entity *TestEntity, size uint64, memoryType UcsMemoryType) unsafe.Pointer {
mmapParams := &UcpMmapParams{}
mmapParams.Allocate().SetLength(size).SetMemoryType(memoryType)
result, err := entity.context.MemMap(mmapParams)
if err != nil {
entity.t.Fatalf("Failed to allocate memory %v", err)
}
entity.mem = result
memAttr, _ := entity.mem.Query(UCP_MEM_ATTR_FIELD_ADDRESS)
return memAttr.Address
}
func createSelfEp(entity *TestEntity) {
if entity.selfEp == nil {
ucpAddress, err := entity.worker.GetAddress()
if err != nil {
entity.t.Fatalf("Failed to get address %v", err)
}
entity.selfEp, err = entity.worker.NewEndpoint((&UcpEpParams{}).SetUcpAddress(ucpAddress))
if err != nil {
entity.t.Fatalf("Failed to create endpoint %v", err)
}
}
}
func sendRecv(entity *TestEntity, sendAddress unsafe.Pointer, sendLength uint64, sendMemType UcsMemoryType,
recvAddr unsafe.Pointer, recvLength uint64, recvMemoryType UcsMemoryType) {
createSelfEp(entity)
recvRequest, _ := entity.worker.RecvTagNonBlocking(recvAddr, recvLength, selfEpTag, selfEpTag, (&UcpRequestParams{}).SetMemType(recvMemoryType))
defer recvRequest.Close()
sendReq, _ := entity.selfEp.SendTagNonBlocking(selfEpTag, sendAddress, sendLength, (&UcpRequestParams{}).SetMemType(sendMemType))
defer sendReq.Close()
for recvRequest.GetStatus() == UCS_INPROGRESS {
entity.worker.Progress()
}
}
func memorySet(entity *TestEntity, data []byte) {
memAttr, _ := entity.mem.Query(UCP_MEM_ATTR_FIELD_ADDRESS, UCP_MEM_ATTR_FIELD_LENGTH, UCP_MEM_ATTR_FIELD_MEM_TYPE)
stringBytes := CBytes(data)
defer FreeNativeMemory((stringBytes))
sendRecv(entity, stringBytes, uint64(len(data)), UCS_MEMORY_TYPE_HOST, memAttr.Address, memAttr.Length, memAttr.MemType)
}
func memoryGet(entity *TestEntity) []byte {
memAttr, _ := entity.mem.Query(UCP_MEM_ATTR_FIELD_ADDRESS, UCP_MEM_ATTR_FIELD_LENGTH, UCP_MEM_ATTR_FIELD_MEM_TYPE)
if memAttr.MemType == UCS_MEMORY_TYPE_HOST {
return GoBytes(memAttr.Address, memAttr.Length)
} else {
recvMem := AllocateNativeMemory(memAttr.Length)
sendRecv(entity, memAttr.Address, memAttr.Length, memAttr.MemType, recvMem, memAttr.Length, UCS_MEMORY_TYPE_HOST)
return GoBytes(recvMem, memAttr.Length)
}
}
|
