From 537ba63547f30ff75dec51a0e41e044b33613c93 Mon Sep 17 00:00:00 2001 From: frreiss Date: Wed, 13 Mar 2019 16:19:23 -0700 Subject: [PATCH 1/3] Add function to create serialized ConfigOptions protos --- tensorflow/go/session.go | 22 ++++++++++++++++++++++ tensorflow/go/session_test.go | 3 ++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/tensorflow/go/session.go b/tensorflow/go/session.go index 48909ffe39e..e1ea3d3d300 100644 --- a/tensorflow/go/session.go +++ b/tensorflow/go/session.go @@ -18,6 +18,7 @@ package tensorflow // #include // #include "tensorflow/c/c_api.h" +// #include "tensorflow/c/c_api_experimental.h" import "C" import ( @@ -349,6 +350,27 @@ func (o *SessionOptions) c() (ret *C.TF_SessionOptions, done func(), err error) }, nil } +// NewConfigOptions generates a serialized ConfigOptions protobuf for use +// in the `Config` field of a `SessionOptions` struct. The function only supports +// the session options available via TensorFlow's experimental C API function +// `TF_CreateConfig()`. THIS API IS UNSTABLE, and the set of available options +// will change in future versions of TensorFlow. +func NewConfigOptions(enableXLACompilation bool, gpuMemoryAllowGrowth bool, numCPUDevices uint) []byte { + // C API expects unsigned chars + enableXLACompilationAsChar := C.uchar(0) + if enableXLACompilation { + enableXLACompilationAsChar = 1 + } + gpuMemoryAllowGrowthAsChar := C.uchar(0) + if gpuMemoryAllowGrowth { + gpuMemoryAllowGrowthAsChar = 1 + } + buf := C.TF_CreateConfig(enableXLACompilationAsChar, gpuMemoryAllowGrowthAsChar, C.uint(numCPUDevices)) + defer C.TF_DeleteBuffer(buf) + // Copy out of C memory. + return C.GoBytes(unsafe.Pointer(buf.data), C.int(buf.length)) +} + // cRunArgs translates the arguments to Session.Run and PartialRun.Run into // values suitable for C library calls. type cRunArgs struct { diff --git a/tensorflow/go/session_test.go b/tensorflow/go/session_test.go index c9bda001671..642d3bb36cf 100644 --- a/tensorflow/go/session_test.go +++ b/tensorflow/go/session_test.go @@ -270,7 +270,8 @@ func TestSessionConfig(t *testing.T) { if err != nil { t.Fatal(err) } - opts := SessionOptions{Config: []byte("(\x01")} + config := NewConfigOptions(true, true, 1) + opts := SessionOptions{Config: config} s, err := NewSession(graph, &opts) if err != nil { t.Fatal(err) From 3b3397c6094fbf7421ef95441b11bf6ccfab6621 Mon Sep 17 00:00:00 2001 From: frreiss Date: Fri, 15 Mar 2019 12:26:02 -0700 Subject: [PATCH 2/3] Introduce Config struct to encapsulate session config options --- tensorflow/go/session.go | 74 +++++++++++++++++++++++++++++------ tensorflow/go/session_test.go | 21 ++-------- 2 files changed, 67 insertions(+), 28 deletions(-) diff --git a/tensorflow/go/session.go b/tensorflow/go/session.go index e1ea3d3d300..c60171af078 100644 --- a/tensorflow/go/session.go +++ b/tensorflow/go/session.go @@ -350,22 +350,74 @@ func (o *SessionOptions) c() (ret *C.TF_SessionOptions, done func(), err error) }, nil } -// NewConfigOptions generates a serialized ConfigOptions protobuf for use -// in the `Config` field of a `SessionOptions` struct. The function only supports -// the session options available via TensorFlow's experimental C API function -// `TF_CreateConfig()`. THIS API IS UNSTABLE, and the set of available options -// will change in future versions of TensorFlow. -func NewConfigOptions(enableXLACompilation bool, gpuMemoryAllowGrowth bool, numCPUDevices uint) []byte { - // C API expects unsigned chars +// JitLevel represents the level of optimization that the XLA compiler +// performs during just-in-time compilation. +type JitLevel int + +const ( + // DEFAULT is the default setting for this version of TensorFlow, + // Currently the default is OFF, but it will change to ON in a future + // version. + DEFAULT JitLevel = 0 + // OFF disables just-in-time compilation and will continue to do so + // even after JIT compilation is enabled by default. + OFF JitLevel = -1 + // ON is a synonym for the ON_1 optimization level. + ON JitLevel = 1 +) + +// Config represents session parameters as encoded in the tensorflow.ConfigProto +// protocol buffer message. +type Config struct { + // GlobalJitLevel controls the degree of optimization that the XLA just-in-time + // compiler will perform. The default is currently "off", but it is expected + // to change to "on" in a future version of TensorFlow. + GlobalJitLevel JitLevel + + // AllowGPUMemoryGrowth controls whether the TensorFlow memory allocator + // pre-allocates the entire specified GPU memory region or instead starts + // with a small block of GPU memory and grows its memory usage as needed. + AllowGPUMemoryGrowth bool + + // NumCPUs is the maximum number of CPU devices that the session will use. + // A value of 0 means "let the system pick an appropriate number" + NumCPUs int + + // This struct only exposes the session options available via TensorFlow's + // experimental C API function `TF_CreateConfig()`. + // TODO(frreiss): Add additional options here as more session options are + // exposed via the C API. +} + +// Bytes generates a serialized ConfigOptions protobuf for use in the `Config` +// field of a `SessionOptions` struct. +func (c *Config) Bytes() []byte { + // The C API expects an unsigned char that is 0 if XLA compilation is off and + // nonzero otherwise. + // There is currently no way in the C API to specify "use TensorFlow's default + // JIT level". The translation logic here ensures that the zero value of + // c.GlobalJitLevel means the same as the default value of + // OptimizerOptions.global_jit_level in the Python API. enableXLACompilationAsChar := C.uchar(0) - if enableXLACompilation { - enableXLACompilationAsChar = 1 + switch c.GlobalJitLevel { + case DEFAULT: + // TODO(frreiss): When the semantics of GlobalJitLevel.DEFAULT change to + // "on", uncomment the following line. + // enableXLACompilationAsChar = C.uchar(1) + case ON: + enableXLACompilationAsChar = C.uchar(1) } gpuMemoryAllowGrowthAsChar := C.uchar(0) - if gpuMemoryAllowGrowth { + if c.AllowGPUMemoryGrowth { gpuMemoryAllowGrowthAsChar = 1 } - buf := C.TF_CreateConfig(enableXLACompilationAsChar, gpuMemoryAllowGrowthAsChar, C.uint(numCPUDevices)) + // The C API doesn't currently have a way to say "let the system pick how many + // CPUs to use," so detect the number of CPUs here. + numCPUDevicesAsUint := C.uint(runtime.NumCPU()) + if c.NumCPUs > 0 { + numCPUDevicesAsUint = C.uint(c.NumCPUs) + } + buf := C.TF_CreateConfig(enableXLACompilationAsChar, gpuMemoryAllowGrowthAsChar, numCPUDevicesAsUint) defer C.TF_DeleteBuffer(buf) // Copy out of C memory. return C.GoBytes(unsafe.Pointer(buf.data), C.int(buf.length)) diff --git a/tensorflow/go/session_test.go b/tensorflow/go/session_test.go index 642d3bb36cf..510108de1ae 100644 --- a/tensorflow/go/session_test.go +++ b/tensorflow/go/session_test.go @@ -250,28 +250,15 @@ func ExamplePartialRun() { } func TestSessionConfig(t *testing.T) { - // Exercise SessionOptions. - // Arguably, a better API would be for SessionOptions.Config to be the - // type generated by the protocol buffer compiler. But for now, the - // tensorflow package continues to be independent of protocol buffers - // and this test exercises the option since the implementation has a - // nuanced conversion to C types. - // - // Till then, the []byte form of Config here was generated using a toy - // tensorflow Python program: - /* - import tensorflow - c = tensorflow.ConfigProto() - c.intra_op_parallelism_threads = 1 - print c.SerializeToString() - */ + // Exercise SessionOptions and Config structs graph := NewGraph() c, err := Const(graph, "Const", int32(14)) if err != nil { t.Fatal(err) } - config := NewConfigOptions(true, true, 1) - opts := SessionOptions{Config: config} + // Use the zero values for Config.GlobalJitLevel and NumCPUs + config := Config{AllowGPUMemoryGrowth: true} + opts := SessionOptions{Config: config.Bytes()} s, err := NewSession(graph, &opts) if err != nil { t.Fatal(err) From aaec622b8a28fcdfbb3847a47e123bdff1529f6b Mon Sep 17 00:00:00 2001 From: frreiss Date: Mon, 18 Mar 2019 13:07:22 -0700 Subject: [PATCH 3/3] Address review comments and add some more API docs --- tensorflow/go/session.go | 29 ++++++++++++++++++----------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/tensorflow/go/session.go b/tensorflow/go/session.go index c60171af078..fc0716bb70f 100644 --- a/tensorflow/go/session.go +++ b/tensorflow/go/session.go @@ -316,6 +316,11 @@ type SessionOptions struct { // Config is a binary-serialized representation of the // tensorflow.ConfigProto protocol message // (https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto). + // You can populate this field in three ways. You can use the zero value + // of this field, which configures the session with a default set of + // options. Or you create a `Config` struct with your options and call that + // struct's `Bytes()` method. Or you can generate a byte string outside of + // Go and paste that string into your Go program as a string literal. Config []byte } @@ -355,15 +360,17 @@ func (o *SessionOptions) c() (ret *C.TF_SessionOptions, done func(), err error) type JitLevel int const ( - // DEFAULT is the default setting for this version of TensorFlow, - // Currently the default is OFF, but it will change to ON in a future - // version. - DEFAULT JitLevel = 0 - // OFF disables just-in-time compilation and will continue to do so + // JitDefault is the default setting for this version of TensorFlow + // and corresponds to the "DEFAULT" level in the Python API. + // Currently the default is `JitOff`, but it will change to `JitOn` in a + // future version. + JitDefault JitLevel = 0 + // JitOff disables just-in-time compilation and will continue to do so // even after JIT compilation is enabled by default. - OFF JitLevel = -1 - // ON is a synonym for the ON_1 optimization level. - ON JitLevel = 1 + JitOff JitLevel = -1 + // JitOn enables just-in-time compilation. It is a synonym for the + // "ON_1" optimization level in the Python API. + JitOn JitLevel = 1 ) // Config represents session parameters as encoded in the tensorflow.ConfigProto @@ -395,16 +402,16 @@ func (c *Config) Bytes() []byte { // The C API expects an unsigned char that is 0 if XLA compilation is off and // nonzero otherwise. // There is currently no way in the C API to specify "use TensorFlow's default - // JIT level". The translation logic here ensures that the zero value of + // JIT level". The translation logic here ensures that the zero value of // c.GlobalJitLevel means the same as the default value of // OptimizerOptions.global_jit_level in the Python API. enableXLACompilationAsChar := C.uchar(0) switch c.GlobalJitLevel { - case DEFAULT: + case JitDefault: // TODO(frreiss): When the semantics of GlobalJitLevel.DEFAULT change to // "on", uncomment the following line. // enableXLACompilationAsChar = C.uchar(1) - case ON: + case JitOn: enableXLACompilationAsChar = C.uchar(1) } gpuMemoryAllowGrowthAsChar := C.uchar(0)