1/* 2Copyright 2017 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 8http://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 internal 18 19/* 20#include <stdlib.h> 21#include <string.h> 22 23#include "tensorflow/c/c_api.h" 24*/ 25import "C" 26 27import ( 28 "errors" 29 "fmt" 30 "runtime" 31 "unsafe" 32 33 "github.com/golang/protobuf/proto" 34 adpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/api_def_go_proto" 35 odpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/op_def_go_proto" 36) 37 38// Encapsulates a collection of API definitions. 39// 40// apiDefMap represents a map from operation name to corresponding 41// ApiDef proto (see 42// https://www.tensorflow.org/code/tensorflow/core/framework/api_def.proto 43// for ApiDef proto definition). 44type apiDefMap struct { 45 c *C.TF_ApiDefMap 46} 47 48// Creates and returns a new apiDefMap instance. 49// 50// oplist is and OpList proto instance (see 51// https://www.tensorflow.org/code/tensorflow/core/framework/op_def.proto 52// for OpList proto definition). 53 54func newAPIDefMap(oplist *odpb.OpList) (*apiDefMap, error) { 55 // Create a buffer containing the serialized OpList. 56 opdefSerialized, err := proto.Marshal(oplist) 57 if err != nil { 58 return nil, fmt.Errorf("could not serialize OpDef for %s", oplist.String()) 59 } 60 data := C.CBytes(opdefSerialized) 61 defer C.free(data) 62 63 opbuf := C.TF_NewBuffer() 64 defer C.TF_DeleteBuffer(opbuf) 65 opbuf.data = data 66 opbuf.length = C.size_t(len(opdefSerialized)) 67 68 // Create ApiDefMap. 69 status := C.TF_NewStatus() 70 defer C.TF_DeleteStatus(status) 71 capimap := C.TF_NewApiDefMap(opbuf, status) 72 if C.TF_GetCode(status) != C.TF_OK { 73 return nil, errors.New(C.GoString(C.TF_Message(status))) 74 } 75 apimap := &apiDefMap{capimap} 76 runtime.SetFinalizer( 77 apimap, 78 func(a *apiDefMap) { 79 C.TF_DeleteApiDefMap(a.c) 80 }) 81 return apimap, nil 82} 83 84// Updates apiDefMap with the overrides specified in `data`. 85// 86// data - ApiDef text proto. 87func (m *apiDefMap) Put(data string) error { 88 cdata := C.CString(data) 89 defer C.free(unsafe.Pointer(cdata)) 90 status := C.TF_NewStatus() 91 defer C.TF_DeleteStatus(status) 92 C.TF_ApiDefMapPut(m.c, cdata, C.size_t(len(data)), status) 93 if C.TF_GetCode(status) != C.TF_OK { 94 return errors.New(C.GoString(C.TF_Message(status))) 95 } 96 return nil 97} 98 99// Returns ApiDef proto instance for the TensorFlow operation 100// named `opname`. 101func (m *apiDefMap) Get(opname string) (*adpb.ApiDef, error) { 102 cname := C.CString(opname) 103 defer C.free(unsafe.Pointer(cname)) 104 status := C.TF_NewStatus() 105 defer C.TF_DeleteStatus(status) 106 apidefBuf := C.TF_ApiDefMapGet( 107 m.c, cname, C.size_t(len(opname)), status) 108 defer C.TF_DeleteBuffer(apidefBuf) 109 if C.TF_GetCode(status) != C.TF_OK { 110 return nil, errors.New(C.GoString(C.TF_Message(status))) 111 } 112 if apidefBuf == nil { 113 return nil, fmt.Errorf("could not find ApiDef for %s", opname) 114 } 115 116 var ( 117 apidef = new(adpb.ApiDef) 118 size = int(apidefBuf.length) 119 // A []byte backed by C memory. 120 // See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices 121 data = (*[1 << 30]byte)(unsafe.Pointer(apidefBuf.data))[:size:size] 122 err = proto.Unmarshal(data, apidef) 123 ) 124 if err != nil { 125 return nil, err 126 } 127 return apidef, nil 128} 129