Optional bitlength when exporting + formatting (#74)

* make max_bitlen optional

* formatting

* separate methods for normal and limbs tensors

* rename to `import_` and `export_`

* update readme example, add it to examples/, rename old example

Co-authored-by: jvmancuso <jason@manc.us>

.gitignore

@@ -5,6 +5,7 @@ bazel*
 **.pyc
 .mypy*
 __pycache__
+**/*egg-info
 
 artifacts/*
 wheelhouse/*

Makefile

@@ -24,9 +24,12 @@ pytest:
 
 fmt:
 	cd tf_big && find . -iname *.h -o -iname *.cc | xargs clang-format -i -style=google
+	isort --atomic --recursive tf_big examples
+	black tf_big examples
 
 lint:
 	cd tf_big && find . -iname *.h -o -iname *.cc | xargs cpplint --filter=-legal/copyright
+	flake8 tf_big examples
 
 download-wheels:
 	rm -rf $(DIR_WHEEL)

README.md

@@ -11,18 +11,17 @@ import tensorflow as tf
 import tf_big
 
 # load large values as strings
-x = tf_big.constant(["100000000000000000000", "200000000000000000000"])
+x = tf_big.constant([["100000000000000000000", "200000000000000000000"]])
 
 # load ordinary TensorFlow tensors
-y = tf_big.convert_to_tensor(tf.constant([3, 4]))
+y = tf_big.import_tensor(tf.constant([[3, 4]]))
 
 # perform computation as usual
 z = x * y
 
-# use TensorFlow sessions to evalutate the results
-with tf.Session() as sess:
-  res = sess.run(z)
-  print(res)
+# export result back into a TensorFlow tensor
+tf_res = tf_big.export_tensor(z)
+print(tf_res)
 ```
 
 ## Installation

examples/simple.py → examples/graph_mode.py

@@ -1,9 +1,10 @@
 import tensorflow as tf
+
 import tf_big
 
 with tf.compat.v1.Session() as sess:
-    x = tf_big.constant([[1,2,3,4]])
-    y = tf_big.constant([[1,2,3,4]])
+    x = tf_big.constant([[1, 2, 3, 4]])
+    y = tf_big.constant([[1, 2, 3, 4]])
     z = x + y
 
     res = sess.run(z)

examples/intro.py

@@ -0,0 +1,15 @@
+import tensorflow as tf
+import tf_big
+
+# load large values as strings
+x = tf_big.constant([["100000000000000000000", "200000000000000000000"]])
+
+# load ordinary TensorFlow tensors
+y = tf_big.import_tensor(tf.constant([[3, 4]]))
+
+# perform computation as usual
+z = x * y
+
+# export result back into a TensorFlow tensor
+tf_res = tf_big.export_tensor(z)
+print(tf_res)

requirements-dev.txt

@@ -1,5 +1,10 @@
 absl-py==0.9.0
+black==19.10b
 cpplint==1.4.4
+flake8-black==0.1.1
+flake8-isort==3.0.0
+flake8==3.7.9
+isort==4.3.21
 numpy==1.16.4
 pip==20.1.1
 setuptools==41.2.0

setup.cfg

@@ -0,0 +1,11 @@
+[tool:isort]
+line_length=88
+force_single_line=True
+
+[flake8]
+max-line-length=88
+extend-ignore=
+    E203  # okay: black
+    T484  # TODO: should not be ignored
+    T499  # TODO: status only?
+    D10,D20,D40  # TODO

tf_big/__init__.py

@@ -1,41 +1,37 @@
-from tf_big.python.tensor import set_secure_default
-from tf_big.python.tensor import get_secure_default
-
 from tf_big.python.tensor import Tensor
-
-from tf_big.python.tensor import constant
-from tf_big.python.tensor import convert_from_tensor
-from tf_big.python.tensor import convert_to_tensor
-
-from tf_big.python.tensor import random_uniform
-from tf_big.python.tensor import random_rsa_modulus
-
 from tf_big.python.tensor import add
-from tf_big.python.tensor import sub
-from tf_big.python.tensor import mul
-from tf_big.python.tensor import pow
+from tf_big.python.tensor import constant
+from tf_big.python.tensor import export_limbs_tensor
+from tf_big.python.tensor import export_tensor
+from tf_big.python.tensor import get_secure_default
+from tf_big.python.tensor import import_limbs_tensor
+from tf_big.python.tensor import import_tensor
+from tf_big.python.tensor import inv
 from tf_big.python.tensor import matmul
 from tf_big.python.tensor import mod
-from tf_big.python.tensor import inv
+from tf_big.python.tensor import mul
+from tf_big.python.tensor import pow
+from tf_big.python.tensor import random_rsa_modulus
+from tf_big.python.tensor import random_uniform
+from tf_big.python.tensor import set_secure_default
+from tf_big.python.tensor import sub
 
 __all__ = [
-  'set_secure_default',
-  'get_secure_default',
-
-  'Tensor',
-
-  'constant',
-  'convert_from_tensor',
-  'convert_to_tensor',
-
-  'random_uniform',
-  'randon_rsa_modulus',
-
-  'add',
-  'sub',
-  'mul',
-  'pow',
-  'matmul',
-  'mod',
-  'inv',
+    "set_secure_default",
+    "get_secure_default",
+    "Tensor",
+    "constant",
+    "export_limbs_tensor",
+    "export_tensor",
+    "import_limbs_tensor",
+    "import_tensor",
+    "random_uniform",
+    "randon_rsa_modulus",
+    "add",
+    "sub",
+    "mul",
+    "pow",
+    "matmul",
+    "mod",
+    "inv",
 ]

tf_big/cc/big_tensor.h

@@ -1,14 +1,13 @@
 #ifndef TF_BIG_CC_BIG_TENSOR_H_
 #define TF_BIG_CC_BIG_TENSOR_H_
 
+#include <fcntl.h>
 #include <gmp.h>
 #include <gmpxx.h>
+#include <unistd.h>
 
 #include <string>
 
-#include <fcntl.h>
-#include <unistd.h>
-
 #include "Eigen/Core"
 #include "Eigen/Dense"
 #include "tensorflow/core/framework/op.h"

tf_big/cc/kernels/big_kernels.cc

@@ -100,55 +100,72 @@ class BigExportLimbsOp : public OpKernel {
       : OpKernel(context) {}
 
   void Compute(OpKernelContext* ctx) override {
-    const Tensor& maxval_tensor = ctx->input(0);
-    auto max_bitlen = maxval_tensor.flat<int32>()(0);
+    const Tensor& maxval_tensor = ctx->input(1);
+    int32_t max_bitlen = maxval_tensor.flat<int32>()(0);
 
-    const BigTensor* curBigTensor = nullptr;
-    TensorShape input_shape = ctx->input(1).shape();
-    OP_REQUIRES_OK(ctx, GetBigTensor(ctx, 1, &curBigTensor));
+    const BigTensor* input = nullptr;
+    TensorShape input_shape = ctx->input(0).shape();
+    OP_REQUIRES_OK(ctx, GetBigTensor(ctx, 0, &input));
+
+    // Compute maxval if left unspecified by user
+    if (max_bitlen < 0) {
+      for (int i = 0; i < input->rows(); i++) {
+        for (int j = 0; j < input->cols(); j++) {
+          auto ele = input->value(i, j).get_mpz_t();
+          int32_t ele_bitlen = mpz_sizeinbase(ele, 2);
+          if (max_bitlen < ele_bitlen) {
+            max_bitlen = ele_bitlen;
+          }
+        }
+      }
+    }
+    OP_REQUIRES(ctx, max_bitlen >= 0,
+                errors::Internal("Malformed max bitlength: ", max_bitlen));
+    unsigned int max_bytelen = max_bitlen * 8;
 
-    Tensor* output;
+    unsigned int header_bytelen = 4;
+    unsigned int header_bitlen = header_bytelen * 8;
+    unsigned int entry_bitlen = header_bitlen + max_bitlen;
+    unsigned int type_bitlen = sizeof(T) * 8;
+    unsigned int num_limbs = (entry_bitlen + type_bitlen - 1) / type_bitlen;
+    unsigned int entry_bytelen = num_limbs * sizeof(T);
 
     TensorShape output_shape;
     output_shape.AddDim(input_shape.dim_size(0));
     output_shape.AddDim(input_shape.dim_size(1));
+    output_shape.AddDim(num_limbs);
 
-    unsigned int type_bitlen = sizeof(T) * 8;
-    unsigned int len_field_bits = 32;
-    unsigned int num_max_limbs =
-        (len_field_bits + max_bitlen + type_bitlen - 1) / type_bitlen;
-
-    output_shape.AddDim(num_max_limbs);
-
+    Tensor* output;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output));
+    auto output_flattened = output->flat<T>();
+    uint8_t* output_data = reinterpret_cast<uint8_t*>(output_flattened.data());
 
-    auto rows = curBigTensor->value.rows();
-    auto cols = curBigTensor->value.cols();
-
-    auto flatened = output->flat<T>();
-    uint8_t* result = reinterpret_cast<uint8_t*>(flatened.data());
-
-    size_t expansion_factor = output->dim_size(2) * sizeof(T);
-    size_t header_len = 4;
     size_t pointer = 0;
+    for (int i = 0; i < input->rows(); i++) {
+      for (int j = 0; j < input->cols(); j++) {
+        auto ele = input->value(i, j).get_mpz_t();
 
-    for (int i = 0; i < rows; i++) {
-      for (int j = 0; j < cols; j++) {
-        unsigned int num =
-            mpz_sizeinbase(curBigTensor->value(i, j).get_mpz_t(), 256);
-        tf_big::encode_length(result + pointer, num);
-
-        size_t exported_len;
-        mpz_export(result + pointer + header_len, &exported_len, 1,
-                   sizeof(uint8_t), 0, 0,
-                   curBigTensor->value(i, j).get_mpz_t());
+        // Write header to output buffer
+        unsigned int ele_bytelen = mpz_sizeinbase(ele, 256);
+        tf_big::encode_length(output_data + pointer, ele_bytelen);
         OP_REQUIRES(
-            ctx, header_len + exported_len <= expansion_factor,
+            ctx, ele_bytelen <= max_bytelen,
             errors::Internal("User selected wrong byte length, required: ",
-                             exported_len, " bytes"));
-        for (size_t k = header_len + exported_len; k < expansion_factor; k++)
-          result[pointer + k] = 0;
-        pointer += expansion_factor;
+                             ele_bytelen, " bytes"));
+
+        // Write element to output buffer
+        size_t exported_bytelen;
+        mpz_export(output_data + pointer + header_bytelen, &exported_bytelen, 1,
+                   sizeof(uint8_t), 0, 0, ele);
+
+        // Zero-out remaining bytes of entry
+        for (size_t k = header_bytelen + exported_bytelen; k < entry_bytelen;
+             k++) {
+          output_data[pointer + k] = 0;
+        }
+
+        // Advance pointer to next entry
+        pointer += entry_bytelen;
       }
     }
   }

tf_big/cc/ops/big_ops.cc

@@ -39,18 +39,17 @@ REGISTER_OP("BigExport")
 
 REGISTER_OP("BigExportLimbs")
     .Attr("dtype: {int32, uint8}")
-    .Input("maxval: int32")
-    .Input("in: variant")
+    .Input("val: variant")
+    .Input("max_bitlen: int32")
     .Output("out: dtype")
     .SetIsStateful()
     .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
-      ::tensorflow::shape_inference::ShapeHandle maxval_shape = c->input(0);
-      TF_RETURN_IF_ERROR(c->WithRank(maxval_shape, 0, &maxval_shape));
-
-      ::tensorflow::shape_inference::ShapeHandle input_shape = c->input(1);
-
+      ::tensorflow::shape_inference::ShapeHandle input_shape = c->input(0);
       TF_RETURN_IF_ERROR(c->WithRank(input_shape, 2, &input_shape));
 
+      ::tensorflow::shape_inference::ShapeHandle max_bitlen_shape = c->input(1);
+      TF_RETURN_IF_ERROR(c->WithRank(max_bitlen_shape, 0, &max_bitlen_shape));
+
       ::tensorflow::shape_inference::ShapeHandle expansion_shape =
           c->MakeShape({c->UnknownDim()});
       ::tensorflow::shape_inference::ShapeHandle out_shape;

tf_big/python/ops/big_ops.py

@@ -2,7 +2,6 @@
 from tensorflow.python.framework.errors import NotFoundError
 from tensorflow.python.platform import resource_loader
 
-
 big_ops_libfile = resource_loader.get_path_to_datafile("_big_ops.so")
 big_ops = tf.load_op_library(big_ops_libfile)