github.com/rohankumardubey/aresdb@v0.0.2-0.20190517170215-e54e3ca06b9c/memutils/memory_test.go (about)

     1  //  Copyright (c) 2017-2018 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package memutils
    16  
    17  import (
    18  	"unsafe"
    19  
    20  	"fmt"
    21  	"github.com/onsi/ginkgo"
    22  	. "github.com/onsi/gomega"
    23  )
    24  
    25  var _ = ginkgo.Describe("memory utils", func() {
    26  
    27  	ginkgo.It("MakeSliceFromCPtr should work", func() {
    28  		// Allocate 5 ints.
    29  		p := HostAlloc(5)
    30  		b := uintptr(p)
    31  		Ω(*(*int)(unsafe.Pointer(b))).Should(BeEquivalentTo(0))
    32  		s := MakeSliceFromCPtr(b, 5)
    33  		s[0] = 1
    34  		Ω(*(*byte)(unsafe.Pointer(b))).Should(BeEquivalentTo(1))
    35  		s[1] = 3
    36  		Ω(*(*byte)(unsafe.Pointer(b + 1))).Should(BeEquivalentTo(3))
    37  		s[4] = 5
    38  		Ω(*(*byte)(unsafe.Pointer(b + 4))).Should(BeEquivalentTo(5))
    39  
    40  		//// Index should be out of bound.
    41  		Ω(func() { fmt.Println(s[5]) }).Should(Panic())
    42  		HostFree(p)
    43  	})
    44  
    45  	ginkgo.It("CreateCudaStream should work", func() {
    46  		Ω(CreateCudaStream(0)).Should(BeZero())
    47  		Ω(CreateCudaStream(10)).Should(BeZero())
    48  	})
    49  
    50  	ginkgo.It("DestroyCudaStream should work", func() {
    51  		Ω(func() { DestroyCudaStream(unsafe.Pointer(nil), 0) }).ShouldNot(Panic())
    52  		Ω(func() { DestroyCudaStream(unsafe.Pointer(nil), 10) }).ShouldNot(Panic())
    53  	})
    54  
    55  	ginkgo.It("DeviceAllocate and DeviceFree should work", func() {
    56  		p1 := DeviceAllocate(0, 0)
    57  		Ω(p1).ShouldNot(BeZero())
    58  		Ω(func() { DeviceFree(p1, 0) }).ShouldNot(Panic())
    59  
    60  		p2 := DeviceAllocate(20, 0)
    61  		Ω(p2).ShouldNot(BeZero())
    62  		Ω(func() { DeviceFree(p2, 0) }).ShouldNot(Panic())
    63  	})
    64  
    65  	ginkgo.It("CudaProfilerStart and CudaProfilerStop should work", func() {
    66  		Ω(func() { CudaProfilerStart() }).ShouldNot(Panic())
    67  		Ω(func() { CudaProfilerStop() }).ShouldNot(Panic())
    68  	})
    69  
    70  	ginkgo.It("GetDeviceCount, GetDeviceMemoryInfo and GetDeviceGlobalMemoryInMB should work", func() {
    71  		deviceCount := GetDeviceCount()
    72  		Ω(deviceCount).Should(BeNumerically(">", 0))
    73  		for device := 0; device < deviceCount; device++ {
    74  			Ω(GetDeviceGlobalMemoryInMB(device)).Should(BeNumerically(">", 0))
    75  			if IsPooledMemory() {
    76  				free, total := GetDeviceMemoryInfo(device)
    77  				Ω(free).Should(BeNumerically(">", 0))
    78  				Ω(total).Should(BeNumerically(">", 0))
    79  			} else {
    80  				Ω(func() { GetDeviceMemoryInfo(device) }).Should(Panic())
    81  			}
    82  		}
    83  	})
    84  
    85  	ginkgo.It("GetFlags should work", func() {
    86  		Ω(func() { GetFlags() }).ShouldNot(Panic())
    87  		Ω(func() { IsDeviceMemoryImplementation() }).ShouldNot(Panic())
    88  		Ω(func() { IsPooledMemory() }).ShouldNot(Panic())
    89  	})
    90  
    91  	ginkgo.It("Init should work", func() {
    92  		Ω(func() { Init() }).ShouldNot(Panic())
    93  	})
    94  
    95  	ginkgo.It("CudaMemCopies should work", func() {
    96  		srcHost := [10]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
    97  		dstHost := [10]byte{}
    98  
    99  		deviceMem1 := DeviceAllocate(10, 0)
   100  		defer DeviceFree(deviceMem1, 0)
   101  		deviceMem2 := DeviceAllocate(10, 0)
   102  		defer DeviceFree(deviceMem2, 0)
   103  
   104  		stream := CreateCudaStream(0)
   105  		AsyncCopyHostToDevice(deviceMem1, unsafe.Pointer(&srcHost[0]), 10, stream, 0)
   106  		AsyncCopyDeviceToDevice(deviceMem2, deviceMem1, 10, stream, 0)
   107  		AsyncCopyDeviceToHost(unsafe.Pointer(&dstHost[0]), deviceMem2, 10, stream, 0)
   108  		Ω(dstHost).Should(Equal(srcHost))
   109  	})
   110  })