<commit_before>#include "gtest/gtest.h"
#include "rmm.h"

// Helper macros to simplify testing for success or failure
#define ASSERT_SUCCESS(res) ASSERT_EQ(RMM_SUCCESS, (res));
#define ASSERT_FAILURE(res) ASSERT_NE(RMM_SUCCESS, (res));

/// Helper class for similar tests
class MemoryManagerTest : public ::testing::Test {
protected:
  virtual void SetUp() {
  	ASSERT_SUCCESS(rmmInitialize());

  	cudaError_t res = cudaStreamCreate(&stream);
  	ASSERT_EQ(cudaSuccess, res);
  }

  virtual void TearDown() {
  	cudaError_t res = cudaStreamDestroy(stream);
  	ASSERT_EQ(cudaSuccess, res);

  	ASSERT_SUCCESS(rmmFinalize());
  }	

  cudaStream_t stream;

  // some useful allocation sizes
  const size_t size_word = 4;
  const size_t size_kb = size_t(1)<<10;
  const size_t size_mb = size_t(1)<<20;
  const size_t size_gb = size_t(1)<<30;
  const size_t size_tb = size_t(1)<<40;
  const size_t size_pb = size_t(1)<<50;
};

// Init / Finalize tests

TEST_F(MemoryManagerTest, Initialize) {
	ASSERT_SUCCESS(rmmInitialize());
}

TEST_F(MemoryManagerTest, Finalize) {
	ASSERT_SUCCESS(rmmFinalize());
}

// zero size tests

TEST_F(MemoryManagerTest, AllocateZeroBytes) {
	char *a = 0;
	ASSERT_SUCCESS(rmmAlloc((void**)&a, 0, stream));
}

TEST_F(MemoryManagerTest, NullPtrAllocateZeroBytes) {
	ASSERT_SUCCESS(rmmAlloc(0, 0, stream));
}

// Bad argument tests

TEST_F(MemoryManagerTest, NullPtrInvalidArgument) {
	rmmError_t res = rmmAlloc(0, 4, stream);
	ASSERT_FAILURE(res);
	ASSERT_EQ(RMM_ERROR_INVALID_ARGUMENT, res);
}

// Simple allocation / free tests

TEST_F(MemoryManagerTest, AllocateWord) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_kb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, AllocateKB) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_word, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, AllocateMB) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_mb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, AllocateGB) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_gb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, AllocateTB) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_tb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, AllocateTooMuch) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_pb, stream);
	ASSERT_FAILURE(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, FreeZero) {
	ASSERT_SUCCESS(rmmFree(0, stream));
}

// Reallocation tests

TEST_F(MemoryManagerTest, ReallocateSmaller) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_mb, stream);
	ASSERT_SUCCESS(res);
	res = rmmRealloc((void**)&a, size_mb / 2, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, ReallocateMuchSmaller) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_gb, stream);
	ASSERT_SUCCESS(res);
	res = rmmRealloc((void**)&a, size_kb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}


TEST_F(MemoryManagerTest, ReallocateLarger) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_mb, stream);
	ASSERT_SUCCESS(res);
	res = rmmRealloc((void**)&a, size_mb * 2, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, ReallocateMuchLarger) {
	char *a = 0;
	rmmError_t res = rmmAlloc((void**)&a, size_kb, stream);
	ASSERT_SUCCESS(res);
	res = rmmRealloc((void**)&a, size_gb, stream);
	ASSERT_SUCCESS(res);
	res = rmmFree(a, stream);
	ASSERT_SUCCESS(res);
}

TEST_F(MemoryManagerTest, GetInfo) {
	size_t freeBefore = 0, totalBefore = 0;
	rmmError_t res = rmmGetInfo(&freeBefore, &totalBefore, stream);
	ASSERT_SUCCESS(res);
	ASSERT_NE(freeBefore, 0);
	ASSERT_NE(totalBefore, 0);

	char *a = 0;
	size_t sz = size_gb / 2;
	res = rmmAlloc((void**)&a, sz, stream);
	ASSERT_SUCCESS(res);

	// make sure the available free memory goes down after an allocation
	size_t freeAfter = 0, totalAfter = 0;
	res = rmmGetInfo(&freeAfter, &totalAfter, stream);
	ASSERT_SUCCESS(res);
	ASSERT_EQ(totalAfter, totalBefore);
	ASSERT_LE(freeAfter, freeBefore);
}<commit_msg>Adding memory manager files and test files missed in previous commit<commit_after>/*
 * Copyright (c) 2018, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 Allocator class compatible with thrust arrays that uses RMM device memory manager.

 Author: <NAME>
 */

#pragma once

#include <thrust/system_error.h>
#include <thrust/system/cuda/error.h>

#include "rmm.h"

template<class T>
class rmm_allocator : public thrust::device_malloc_allocator<T>
{
  public:
    using value_type = T;

    // TODO: make stream part of the state rather than using NULL stream
    //rmm_allocator(cudaStream_t stream) : stream(stream) {}

    typedef thrust::device_ptr<T>  pointer;
    inline pointer allocate(size_t n)
    {
      value_type* result = nullptr;
  
      rmmError_t error = rmmAlloc((void**)&result, n*sizeof(T), 0); //stream);
     
      if(error != RMM_SUCCESS)
      {
        throw thrust::system_error(error, thrust::cuda_category(), "rmm_allocator::allocate(): rmmAlloc");
      }
  
      return thrust::device_pointer_cast(result);
    }
  
    inline void deallocate(pointer ptr, size_t)
    {
      rmmError_t error = rmmFree(thrust::raw_pointer_cast(ptr), 0);//stream);
  
      if(error != RMM_SUCCESS)
      {
        throw thrust::system_error(error, thrust::cuda_category(), "rmm_allocator::deallocate(): rmmFree");
      }
    }

  //private:
  	//cudaStream_t stream;
};
{'fraction_non_alphanumeric': 0.13280881195908734, 'fraction_numerical': 0.00833988985051141, 'mean_word_length': 4.923578751164958, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6324512', 'n_tokens_mistral': 2348, 'n_tokens_neox': 2283, 'n_words': 538}