1/* 2Copyright 2018 The TensorFlow Authors. All Rights Reserved. 3 4Licensed under the Apache License, Version 2.0 (the "License"); 5you may not use this file except in compliance with the License. 6You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10Unless required by applicable law or agreed to in writing, software 11distributed under the License is distributed on an "AS IS" BASIS, 12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13See the License for the specific language governing permissions and 14limitations under the License. 15*/ 16 17package tensorflow 18 19// #include <stdlib.h> 20// #include "tensorflow/c/c_api.h" 21// #include "tensorflow/c/eager/c_api.h" 22import "C" 23import ( 24 "fmt" 25 "runtime" 26) 27 28// ContextOptions contains configuration information for a session 29type ContextOptions struct { 30 // Config is a binary-serialized representation of the 31 // tensorflow.ConfigProto protocol message 32 // (https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto). 33 Config []byte 34 35 // Sets the default execution mode 36 Async bool 37} 38 39// c converts the ContextOptions to the C API's TF_ContextOptions. 40// Caller takes ownership of returned object. 41func (o *ContextOptions) c() (*C.TFE_ContextOptions, error) { 42 opt := C.TFE_NewContextOptions() 43 if o == nil { 44 return opt, nil 45 } 46 47 if sz := len(o.Config); sz > 0 { 48 status := newStatus() 49 cConfig := C.CBytes(o.Config) 50 C.TFE_ContextOptionsSetConfig(opt, cConfig, C.size_t(sz), status.c) 51 C.free(cConfig) 52 if err := status.Err(); err != nil { 53 C.TFE_DeleteContextOptions(opt) 54 return nil, fmt.Errorf("invalid ContextOptions.Config: %v", err) 55 } 56 } 57 58 var async uint8 59 if o.Async { 60 async = 1 61 } 62 C.TFE_ContextOptionsSetAsync(opt, C.uchar(async)) 63 64 return opt, nil 65} 66 67// Context for executing operations eagerly. 68// 69// A Context allows operations to be executed immediately. It encapsulates 70// information such as the available devices, resource manager etc. It also 71// allows the user to configure execution using a ConfigProto, as they can 72// configure a Session when executing a Graph. 73type Context struct { 74 c *C.TFE_Context 75} 76 77// NewContext creates a new context for eager execution. 78// options may be nil to use the default options. 79func NewContext(options *ContextOptions) (*Context, error) { 80 status := newStatus() 81 cOpt, err := options.c() 82 if err != nil { 83 return nil, err 84 } 85 defer C.TFE_DeleteContextOptions(cOpt) 86 cContext := C.TFE_NewContext(cOpt, status.c) 87 if err := status.Err(); err != nil { 88 return nil, err 89 } 90 91 c := &Context{c: cContext} 92 runtime.SetFinalizer(c, (*Context).finalizer) 93 return c, nil 94} 95 96func (c *Context) finalizer() { 97 C.TFE_DeleteContext(c.c) 98} 99 100// ListDevices returns the list of devices associated with a Context. 101func (c *Context) ListDevices() ([]Device, error) { 102 status := newStatus() 103 devicesList := C.TFE_ContextListDevices(c.c, status.c) 104 if err := status.Err(); err != nil { 105 return nil, fmt.Errorf("SessionListDevices() failed: %v", err) 106 } 107 defer C.TF_DeleteDeviceList(devicesList) 108 return deviceSliceFromDeviceList(devicesList) 109} 110