Yu Yang
8 years ago
commit
ff000ae76a
@ -0,0 +1,181 @@
|
|||||||
|
package pserver
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/PaddlePaddle/Paddle/go/utils/networkhelper"
|
||||||
|
"github.com/coreos/etcd/clientv3"
|
||||||
|
"github.com/coreos/etcd/clientv3/concurrency"
|
||||||
|
log "github.com/sirupsen/logrus"
|
||||||
|
)
|
||||||
|
|
||||||
|
// EtcdClient is the etcd client that the pserver uses for fault
|
||||||
|
// tolerance, service registry and coordination.
|
||||||
|
type EtcdClient struct {
|
||||||
|
numPservers int
|
||||||
|
etcdEndpoints string
|
||||||
|
etcdClient *clientv3.Client
|
||||||
|
// etcdTimeout is also used as retry intervals.
|
||||||
|
etcdTimeout time.Duration
|
||||||
|
// FIXME: ensure GetExternalIP gets the correct ip for trainers to connect.
|
||||||
|
externalIP string
|
||||||
|
// desired number of pservers in the job.
|
||||||
|
// assume desired will not change during one training job.
|
||||||
|
desired int
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewEtcdClient creates an EtcdClient
|
||||||
|
func NewEtcdClient(endpoints string, numPservers int, timeout time.Duration) *EtcdClient {
|
||||||
|
return &EtcdClient{
|
||||||
|
etcdTimeout: timeout,
|
||||||
|
numPservers: numPservers,
|
||||||
|
etcdEndpoints: endpoints,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Register registers the pserver on etcd
|
||||||
|
//
|
||||||
|
// Register returns the index of the current pserver.
|
||||||
|
func (e *EtcdClient) Register() (int, error) {
|
||||||
|
|
||||||
|
var err error
|
||||||
|
e.externalIP, err = networkhelper.GetExternalIP()
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// initialize connection to etcd.
|
||||||
|
ep := strings.Split(e.etcdEndpoints, ",")
|
||||||
|
for {
|
||||||
|
cli, err := clientv3.New(clientv3.Config{
|
||||||
|
Endpoints: ep,
|
||||||
|
DialTimeout: e.etcdTimeout,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
log.Errorf("connect to etcd error: %v", err)
|
||||||
|
time.Sleep(e.etcdTimeout)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
e.etcdClient = cli
|
||||||
|
log.Debugf("inited client to %s", e.etcdEndpoints)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
// init /ps_desired using transaction, for multiple pservers may want to write
|
||||||
|
// it at the same time.
|
||||||
|
for {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||||
|
_, err := e.initDesiredPsercers(ctx, e.numPservers)
|
||||||
|
cancel()
|
||||||
|
if err != nil {
|
||||||
|
log.Warn(err)
|
||||||
|
time.Sleep(e.etcdTimeout)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
// TODO: when implementing extending or reducing pservers, /ps_desired is
|
||||||
|
// changed, then we need to watch /ps_desired node for events. For now, just
|
||||||
|
// write once when init and read from it.
|
||||||
|
// wait and set s.desired init value
|
||||||
|
for {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||||
|
resp, err := e.etcdClient.Get(ctx, PsDesired)
|
||||||
|
cancel()
|
||||||
|
if err != nil {
|
||||||
|
log.Errorf("getting %s error: %v", PsDesired, err)
|
||||||
|
time.Sleep(e.etcdTimeout)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if len(resp.Kvs) != 0 {
|
||||||
|
e.desired, err = strconv.Atoi(string(resp.Kvs[0].Value))
|
||||||
|
if err != nil {
|
||||||
|
log.Errorf("value of %s invalid %v\n", PsDesired, err)
|
||||||
|
time.Sleep(e.etcdTimeout)
|
||||||
|
// NOTE: wait util ps_desired value change
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var pserverIdx int
|
||||||
|
// try register pserver node on etcd
|
||||||
|
for {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||||
|
var err error
|
||||||
|
pserverIdx, err = e.registerPserverEtcd(ctx)
|
||||||
|
cancel()
|
||||||
|
if err != nil {
|
||||||
|
log.Warn(err)
|
||||||
|
time.Sleep(e.etcdTimeout)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
return pserverIdx, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *EtcdClient) initDesiredPsercers(ctx context.Context, numPservers int) (*clientv3.TxnResponse, error) {
|
||||||
|
return concurrency.NewSTM(e.etcdClient, func(c concurrency.STM) error {
|
||||||
|
dsStr := c.Get(PsDesired)
|
||||||
|
if dsStr == "" {
|
||||||
|
c.Put(PsDesired, strconv.Itoa(numPservers))
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}, concurrency.WithAbortContext(ctx), concurrency.WithIsolation(concurrency.RepeatableReads))
|
||||||
|
}
|
||||||
|
|
||||||
|
// registerPserverEtcd registers pserver node on etcd using transaction.
|
||||||
|
func (e *EtcdClient) registerPserverEtcd(ctx context.Context) (int, error) {
|
||||||
|
var idx int
|
||||||
|
_, err := concurrency.NewSTM(e.etcdClient, func(c concurrency.STM) error {
|
||||||
|
registered := false
|
||||||
|
for i := 0; i < e.desired; i++ {
|
||||||
|
psKey := "/ps/" + strconv.Itoa(i)
|
||||||
|
log.Debugf("checking %s", psKey)
|
||||||
|
ps := c.Get(psKey)
|
||||||
|
log.Debugf("got value (%s) for key: %s", ps, psKey)
|
||||||
|
|
||||||
|
if ps == "" {
|
||||||
|
resp, err := e.etcdClient.Grant(context.TODO(), 5)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
// find the first id and write info
|
||||||
|
c.Put(psKey, e.externalIP, clientv3.WithLease(resp.ID))
|
||||||
|
log.Debugf("set pserver node %s with value %s", psKey, e.externalIP)
|
||||||
|
ch, kaerr := e.etcdClient.KeepAlive(context.TODO(), resp.ID)
|
||||||
|
if kaerr != nil {
|
||||||
|
log.Errorf("keepalive etcd node error: %v", kaerr)
|
||||||
|
return kaerr
|
||||||
|
}
|
||||||
|
|
||||||
|
// Eat the keep alive message so etcd
|
||||||
|
// will not expire the lease.
|
||||||
|
go func(ch <-chan *clientv3.LeaseKeepAliveResponse) {
|
||||||
|
ka := <-ch
|
||||||
|
log.Debugf("keepalive: %d\n", ka.TTL)
|
||||||
|
}(ch)
|
||||||
|
log.Debug("register finished")
|
||||||
|
idx = i
|
||||||
|
registered = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if registered == true {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return errors.New("not registerd, may due to already have enough pservers")
|
||||||
|
}, concurrency.WithAbortContext(ctx), concurrency.WithIsolation(concurrency.RepeatableReads))
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return idx, nil
|
||||||
|
}
|
||||||
@ -0,0 +1,5 @@
|
|||||||
|
---
|
||||||
|
Language: Cpp
|
||||||
|
BasedOnStyle: Google
|
||||||
|
Standard: Cpp11
|
||||||
|
...
|
||||||
@ -0,0 +1 @@
|
|||||||
|
add_subdirectory(detail)
|
||||||
@ -0,0 +1,7 @@
|
|||||||
|
if(${WITH_GPU})
|
||||||
|
nv_library(system_allocator SRCS system_allocator.cc DEPS gflags)
|
||||||
|
nv_test(system_allocator_test SRCS system_allocator_test.cc DEPS system_allocator gflags)
|
||||||
|
else(${WITH_GPU})
|
||||||
|
cc_library(system_allocator SRCS system_allocator.cc DEPS gflags)
|
||||||
|
cc_test(system_allocator_test SRCS system_allocator_test.cc DEPS system_allocator gflags)
|
||||||
|
endif(${WITH_GPU})
|
||||||
@ -0,0 +1,35 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include "paddle/memory/detail/buddy_allocator.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace memory {
|
||||||
|
namespace detail {
|
||||||
|
|
||||||
|
BuddyAllocator::BuddyAllocator(size_t pool_size, size_t max_pools,
|
||||||
|
SystemAllocator* system_allocator)
|
||||||
|
: pool_size_(pool_size),
|
||||||
|
max_pools_(max_pools),
|
||||||
|
system_allocator_(system_allocator) {
|
||||||
|
PADDLE_ASSERT(pool_size > 0);
|
||||||
|
PADDLE_ASSERT(max_pools > 0);
|
||||||
|
PADDLE_ASSERT(system_allocator != nullptr);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace detail
|
||||||
|
} // namespace memory
|
||||||
|
} // namespace paddle
|
||||||
@ -0,0 +1,86 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include "paddle/memory/detail/system_allocator.h"
|
||||||
|
|
||||||
|
#include <mutex>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace memory {
|
||||||
|
namespace detail {
|
||||||
|
|
||||||
|
class BuddyAllocator {
|
||||||
|
public:
|
||||||
|
BuddyAllocator(size_t pool_size, size_t max_pools,
|
||||||
|
SystemAllocator* system_allocator);
|
||||||
|
~BuddyAllocator();
|
||||||
|
|
||||||
|
void* Alloc(size_t size);
|
||||||
|
void Free(void*);
|
||||||
|
size_t Used();
|
||||||
|
|
||||||
|
private:
|
||||||
|
struct Block {
|
||||||
|
size_t size_;
|
||||||
|
Block* left_; // left buddy
|
||||||
|
Block* right_; // right buddy
|
||||||
|
};
|
||||||
|
|
||||||
|
// Initially, there is only one pool. If a Alloc founds not enough
|
||||||
|
// memory from that pool, and there has not been max_num_pools_,
|
||||||
|
// create a new pool by calling system_allocator_.Alloc(pool_size_).
|
||||||
|
std::vector<void*> pools_;
|
||||||
|
|
||||||
|
size_t pool_size_; // the size of each pool;
|
||||||
|
size_t max_num_pools_; // the size of all pools;
|
||||||
|
|
||||||
|
SystemAllocator* system_allocator_;
|
||||||
|
|
||||||
|
std::mutex mutex_;
|
||||||
|
|
||||||
|
// Disable copy and assignment.
|
||||||
|
BuddyAllocator(const BuddyAllocator&) = delete;
|
||||||
|
BuddyAllocator& operator=(const BuddyAllocator&) = delete;
|
||||||
|
};
|
||||||
|
|
||||||
|
BuddyAllocator<CPUAllocator>* GetCPUBuddyAllocator() {
|
||||||
|
static BuddyAllocator<CPUAllocator>* a = nullptr;
|
||||||
|
if (a == nullptr) {
|
||||||
|
a = new BuddyAllocator<CPUAllocator>();
|
||||||
|
}
|
||||||
|
return a;
|
||||||
|
}
|
||||||
|
|
||||||
|
#ifndef PADDLE_ONLY_CPU // The following code are for CUDA.
|
||||||
|
|
||||||
|
BuddyAllocator<GPUAllocator>* GetGPUBuddyAllocator(int gpu_id) {
|
||||||
|
static BuddyAllocator<GPUAllocator>** as = NULL;
|
||||||
|
if (as == NULL) {
|
||||||
|
int gpu_num = platform::GetDeviceCount();
|
||||||
|
as = new BuddyAllocator<GPUAllocator>*[gpu_num];
|
||||||
|
for (int gpu = 0; gpu < gpu_num; gpu++) {
|
||||||
|
as[gpu] = new BuddyAllocator<GPUAllocator>();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return as[gpu_id];
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
|
||||||
|
} // namespace detail
|
||||||
|
} // namespace memory
|
||||||
|
} // namespace paddle
|
||||||
@ -0,0 +1,90 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#include "paddle/memory/detail/system_allocator.h"
|
||||||
|
|
||||||
|
#include <stdlib.h> // for malloc and free
|
||||||
|
#include <sys/mman.h> // for mlock and munlock
|
||||||
|
|
||||||
|
#include "gflags/gflags.h"
|
||||||
|
#include "paddle/platform/assert.h"
|
||||||
|
#include "paddle/platform/cuda.h"
|
||||||
|
|
||||||
|
// If use_pinned_memory is true, CPUAllocator calls mlock, which
|
||||||
|
// returns pinned and locked memory as staging areas for data exchange
|
||||||
|
// between host and device. Allocates too much would reduce the amount
|
||||||
|
// of memory available to the system for paging. So, by default, we
|
||||||
|
// should set false to use_pinned_memory.
|
||||||
|
DEFINE_bool(use_pinned_memory, false,
|
||||||
|
"If set, allocate cpu/gpu pinned memory.");
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace memory {
|
||||||
|
namespace detail {
|
||||||
|
|
||||||
|
void* CPUAllocator::Alloc(size_t size) {
|
||||||
|
// According to http://www.cplusplus.com/reference/cstdlib/malloc/,
|
||||||
|
// malloc might not return nullptr if size is zero, but the returned
|
||||||
|
// pointer shall not be dereferenced -- so we make it nullptr.
|
||||||
|
if (size <= 0) return nullptr;
|
||||||
|
|
||||||
|
void* p = malloc(size);
|
||||||
|
if (p != nullptr && FLAGS_use_pinned_memory) {
|
||||||
|
mlock(p, size);
|
||||||
|
}
|
||||||
|
return p;
|
||||||
|
}
|
||||||
|
|
||||||
|
void CPUAllocator::Free(void* p, size_t size) {
|
||||||
|
if (p != nullptr && FLAGS_use_pinned_memory) {
|
||||||
|
munlock(p, size);
|
||||||
|
}
|
||||||
|
free(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
|
||||||
|
void* GPUAllocator::Alloc(size_t size) {
|
||||||
|
// CUDA documentation doesn't explain if cudaMalloc returns nullptr
|
||||||
|
// if size is 0. We just make sure it does.
|
||||||
|
if (size <= 0) {
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
|
void* p = 0;
|
||||||
|
cudaError_t result =
|
||||||
|
FLAGS_use_pinned_memory ? cudaMallocHost(&p, size) : cudaMalloc(&p, size);
|
||||||
|
if (result != cudaSuccess) {
|
||||||
|
cudaGetLastError(); // clear error if there is any.
|
||||||
|
}
|
||||||
|
return result == cudaSuccess ? p : nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
|
void GPUAllocator::Free(void* p, size_t size) {
|
||||||
|
// Purposefully allow cudaErrorCudartUnloading, because
|
||||||
|
// that is returned if you ever call cudaFree after the
|
||||||
|
// driver has already shutdown. This happens only if the
|
||||||
|
// process is terminating, in which case we don't care if
|
||||||
|
// cudaFree succeeds.
|
||||||
|
cudaError_t err = FLAGS_use_pinned_memory ? cudaFreeHost(p) : cudaFree(p);
|
||||||
|
if (err != cudaErrorCudartUnloading) {
|
||||||
|
platform::throw_on_error(err, "cudaFree{Host} failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
|
||||||
|
} // namespace detail
|
||||||
|
} // namespace memory
|
||||||
|
} // namespace paddle
|
||||||
@ -0,0 +1,53 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include <stddef.h> // for size_t
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace memory {
|
||||||
|
namespace detail {
|
||||||
|
|
||||||
|
// SystemAllocator is the parent class of CPUAllocator and
|
||||||
|
// GPUAllocator. A BuddyAllocator object uses a SystemAllocator*
|
||||||
|
// pointing to the underlying system allocator. An alternative to
|
||||||
|
// this class hierarchy is to pass a system allocator class to
|
||||||
|
// BuddyAllocator as a template parameter. This approach makes
|
||||||
|
// BuddyAllocator a class template, and it's very complicated
|
||||||
|
// algorithm would make the buddy_allocator.h messy.
|
||||||
|
class SystemAllocator {
|
||||||
|
public:
|
||||||
|
virtual ~SystemAllocator() {}
|
||||||
|
virtual void* Alloc(size_t size) = 0;
|
||||||
|
virtual void Free(void* p, size_t size) = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
class CPUAllocator : public SystemAllocator {
|
||||||
|
public:
|
||||||
|
virtual void* Alloc(size_t size);
|
||||||
|
virtual void Free(void* p, size_t size);
|
||||||
|
};
|
||||||
|
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
class GPUAllocator : public SystemAllocator {
|
||||||
|
public:
|
||||||
|
virtual void* Alloc(size_t size);
|
||||||
|
virtual void Free(void* p, size_t size);
|
||||||
|
};
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
|
||||||
|
} // namespace detail
|
||||||
|
} // namespace memory
|
||||||
|
} // namespace paddle
|
||||||
@ -0,0 +1,71 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#include "paddle/memory/detail/system_allocator.h"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "gflags/gflags.h"
|
||||||
|
#include "gtest/gtest.h"
|
||||||
|
|
||||||
|
DECLARE_bool(use_pinned_memory);
|
||||||
|
|
||||||
|
void TestAllocator(paddle::memory::detail::SystemAllocator& a, size_t size) {
|
||||||
|
bool freed = false;
|
||||||
|
{
|
||||||
|
void* p = a.Alloc(size);
|
||||||
|
if (size > 0) {
|
||||||
|
EXPECT_NE(p, nullptr);
|
||||||
|
} else {
|
||||||
|
EXPECT_EQ(p, nullptr);
|
||||||
|
}
|
||||||
|
|
||||||
|
int* i = static_cast<int*>(p);
|
||||||
|
std::shared_ptr<int> ptr(i, [&](void* p) {
|
||||||
|
freed = true;
|
||||||
|
a.Free(p, size);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
EXPECT_TRUE(freed);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(CPUAllocator, NoLockMem) {
|
||||||
|
FLAGS_use_pinned_memory = false;
|
||||||
|
paddle::memory::detail::CPUAllocator a;
|
||||||
|
TestAllocator(a, 2048);
|
||||||
|
TestAllocator(a, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(CPUAllocator, LockMem) {
|
||||||
|
FLAGS_use_pinned_memory = true;
|
||||||
|
paddle::memory::detail::CPUAllocator a;
|
||||||
|
TestAllocator(a, 2048);
|
||||||
|
TestAllocator(a, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
TEST(GPUAllocator, NoStaging) {
|
||||||
|
FLAGS_use_pinned_memory = false;
|
||||||
|
paddle::memory::detail::GPUAllocator a;
|
||||||
|
TestAllocator(a, 2048);
|
||||||
|
TestAllocator(a, 0);
|
||||||
|
}
|
||||||
|
TEST(GPUAllocator, Staging) {
|
||||||
|
FLAGS_use_pinned_memory = true;
|
||||||
|
paddle::memory::detail::GPUAllocator a;
|
||||||
|
TestAllocator(a, 2048);
|
||||||
|
TestAllocator(a, 0);
|
||||||
|
}
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
@ -0,0 +1,59 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#include "paddle/memory/memory.h"
|
||||||
|
#include "paddle/memory/detail/buddy_allocator.h"
|
||||||
|
#include "paddle/memory/detail/system_allocator.h"
|
||||||
|
#include "paddle/platform/assert.h"
|
||||||
|
|
||||||
|
#include <boost/variant.hpp>
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace memory {
|
||||||
|
|
||||||
|
void* Alloc(platform::Place pl, size_t size) {
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
if (paddle::platform::is_gpu_place(pl)) {
|
||||||
|
size_t gpu_id = boost::get<platform::GPUPlace>(pl).device;
|
||||||
|
return detail::GetGPUBuddyAllocator(gpu_id)->Alloc(size);
|
||||||
|
}
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
PADDLE_ASSERT(paddle::platform::is_cpu_place(pl));
|
||||||
|
return detail::GetCPUBuddyAllocator()->Alloc(size);
|
||||||
|
}
|
||||||
|
|
||||||
|
void Free(paddle::platform::Place pl, void* p) {
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
if (paddle::platform::is_gpu_place(pl)) {
|
||||||
|
size_t gpu_id = boost::get<platform::GPUPlace>(pl).device;
|
||||||
|
detail::GetGPUBuddyAllocator(gpu_id)->Free(p);
|
||||||
|
}
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
PADDLE_ASSERT(paddle::platform::is_cpu_place(pl));
|
||||||
|
detail::GetCPUBuddyAllocator()->Free(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
size_t Used(paddle::platform::Place pl) {
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
if (paddle::platform::is_gpu_place(pl)) {
|
||||||
|
size_t gpu_id = boost::get<platform::GPUPlace>(pl).device;
|
||||||
|
return detail::GetGPUBuddyAllocator(gpu_id)->Used();
|
||||||
|
}
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
|
PADDLE_ASSERT(paddle::platform::is_cpu_place(pl));
|
||||||
|
return detail::GetCPUBuddyAllocator()->Used();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace memory
|
||||||
|
} // namespace paddle
|
||||||
@ -0,0 +1,40 @@
|
|||||||
|
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#ifndef PADDLE_ONLY_CPU
|
||||||
|
|
||||||
|
#include <thrust/system/cuda/error.h>
|
||||||
|
#include <thrust/system_error.h>
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace platform {
|
||||||
|
|
||||||
|
inline void throw_on_error(cudaError_t e, const char* message) {
|
||||||
|
if (e) {
|
||||||
|
throw thrust::system_error(e, thrust::cuda_category(), message);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
int GetDeviceCount(void) {
|
||||||
|
int count;
|
||||||
|
throw_on_error(cudaGetDeviceCount(&count), "cudaGetDeviceCount failed");
|
||||||
|
return count;
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace platform
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
#endif // PADDLE_ONLY_CPU
|
||||||
@ -1,10 +0,0 @@
|
|||||||
#include <gtest/gtest.h>
|
|
||||||
#include <paddle/platform/must_check.h>
|
|
||||||
|
|
||||||
int __must_check SomeFunctionMustCheck() { return 0; }
|
|
||||||
|
|
||||||
TEST(MustCheck, all) {
|
|
||||||
// This line should not be compiled, because the
|
|
||||||
// return value of SomeFunctionMustCheck marked as __must_check
|
|
||||||
// SomeFunctionMustCheck();
|
|
||||||
}
|
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue