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