استخدام تجميع AOT

ما هو tfcompile؟

tfcompile هي أداة مستقلة تعمل في الوقت المحدّد مسبقًا (AOT) على تجميع رسوم TensorFlow البيانية في شكل رمز برمجي تنفيذي. يمكنها تقليل إجمالي الحجم الثنائي وتجنب بعض النفقات العامة في وقت التشغيل. تتمثل حالة استخدام tfcompile النموذجية في تجميع رسم بياني للاستنتاج في رمز برمجي تنفيذي للأجهزة الجوّالة.

يتم تنفيذ الرسم البياني TensorFlow عادةً من خلال وقت تشغيل TensorFlow. يتحمل هذا بعض الوقت الزائد لتنفيذ كل عقدة في الرسم البياني. يؤدي ذلك أيضًا إلى إجمالي حجم ثنائي أكبر، حيث يجب أن تكون التعليمات البرمجية لوقت تشغيل TensorFlow متاحة، بالإضافة إلى الرسم البياني نفسه. لا يستخدم الرمز القابل للتنفيذ الذي يتم إنشاؤه من خلال tfcompile وقت تشغيل TensorFlow، بل يعتمد فقط على النواة المستخدَمة في العمليات الحسابية.

تم إنشاء برنامج التحويل البرمجي على إطار عمل XLA. يتوفّر الرمز البرمجي الذي يربط بين TensorFlow وإطار عمل XLA ضمن tensorflow/compiler.

ما هي وظيفة tfcompile؟

تأخذ tfcompile رسمًا بيانيًا فرعيًا، يتم تحديده من خلال مفاهيم TensorFlow الخاصة بالخلاصات وعمليات الجلب، وتنشئ دالة تنفّذ هذا الرسم البياني الفرعي. feeds هي وسيطات الإدخال للدالة، وfetches هي وسيطات الناتج للدالة. يجب تحديد جميع الإدخالات بالكامل عن طريق الخلاصات، ولا يمكن أن يحتوي الرسم البياني الفرعي الناتج والمقتبس على عنصر نائب أو عُقد متغيرة. من الشائع تحديد جميع العناصر النائبة والمتغيرات كخلاصات، مما يضمن أن الرسم البياني الفرعي الناتج لم يعد يحتوي على هذه العُقد. يتم حزم الدالة التي تم إنشاؤها كـ cc_library، مع ملف عنوان يُصدِّر توقيع الدالة، وملف كائن يحتوي على التنفيذ. يكتب المستخدم تعليمة برمجية لاستدعاء الدالة التي تم إنشاؤها حسب الحاجة.

استخدام tfcompile

يوضّح هذا القسم بالتفصيل الخطوات عالية المستوى لإنشاء برنامج ثنائي قابل للتنفيذ باستخدام tfcompile من رسم بياني فرعي من TensorFlow. الخطوات كالآتي:

• الخطوة 1: تهيئة الرسم البياني الفرعي لتجميعه
• الخطوة 2: استخدام ماكرو إصدار tf_library لتجميع الرسم البياني الفرعي
• الخطوة 3: كتابة التعليمة البرمجية لاستدعاء الرسم البياني الفرعي
• الخطوة 4: إنشاء البرنامج الثنائي النهائي

الخطوة 1: تهيئة الرسم البياني الفرعي لتجميعه

حدد الخلاصات وعمليات الجلب التي تتوافق مع وسيطات الإدخال والإخراج للدالة التي تم إنشاؤها. بعد ذلك، عليك إعداد feeds وfetches في نموذج tensorflow.tf2xla.Config.

# Each feed is a positional input argument for the generated function.  The order
# of each entry matches the order of each input argument.  Here “x_hold” and “y_hold”
# refer to the names of placeholder nodes defined in the graph.
feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}

# Each fetch is a positional output argument for the generated function.  The order
# of each entry matches the order of each output argument.  Here “x_y_prod”
# refers to the name of a matmul node defined in the graph.
fetch {
  id { node_name: "x_y_prod" }
}

الخطوة 2: استخدام وحدة ماكرو إصدار tf_library لتجميع الرسم البياني الفرعي

تحوِّل هذه الخطوة الرسم البياني إلى cc_library باستخدام ماكرو الإصدار tf_library. يتكوّن cc_library من ملف عنصر يحتوي على الرمز الذي تم إنشاؤه من الرسم البياني، إلى جانب ملف رأس يتيح الوصول إلى الرمز الذي تم إنشاؤه. يستخدم tf_library tfcompile لتحويل الرسم البياني TensorFlow إلى رمز برمجي قابل للتنفيذ.

load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# Use the tf_library macro to compile your graph into executable code.
tf_library(
    # name is used to generate the following underlying build rules:
    # <name>           : cc_library packaging the generated header and object files
    # <name>_test      : cc_test containing a simple test and benchmark
    # <name>_benchmark : cc_binary containing a stand-alone benchmark with minimal deps;
    #                    can be run on a mobile device
    name = "test_graph_tfmatmul",
    # cpp_class specifies the name of the generated C++ class, with namespaces allowed.
    # The class will be generated in the given namespace(s), or if no namespaces are
    # given, within the global namespace.
    cpp_class = "foo::bar::MatMulComp",
    # graph is the input GraphDef proto, by default expected in binary format.  To
    # use the text format instead, just use the ‘.pbtxt’ suffix.  A subgraph will be
    # created from this input graph, with feeds as inputs and fetches as outputs.
    # No Placeholder or Variable ops may exist in this subgraph.
    graph = "test_graph_tfmatmul.pb",
    # config is the input Config proto, by default expected in binary format.  To
    # use the text format instead, use the ‘.pbtxt’ suffix.  This is where the
    # feeds and fetches were specified above, in the previous step.
    config = "test_graph_tfmatmul.config.pbtxt",
)

لإنشاء نموذج GraphDef (test_graph_tfmatmul.pb) في هذا المثال، شغِّل make_test_graphs.py وتحديد موقع الإخراج باستخدام العلامة --out_dir.

تحتوي الرسوم البيانية النموذجية على Variables التي تمثل الأوزان التي تم تعلُّمها من خلال التدريب، ولكن لا يمكن لـ tfcompile تجميع رسم بياني فرعي يحتوي على Variables. تحوّل أداة freeze_graph.py المتغيرات إلى ثوابت باستخدام القيم المخزنة في ملف نقطة تفتيش. تيسيرًا للأمر، تتيح وحدة الماكرو tf_library الوسيطة freeze_checkpoint التي تشغّل الأداة. لمزيد من الأمثلة، راجع tensorflow/compiler/aot/tests/BUILD.

يتم تجميع الثوابت التي تظهر في الرسم البياني الفرعي المجمَّع مباشرةً في التعليمة البرمجية التي تم إنشاؤها. لتمرير الثوابت في الدالة المُنشأة، بدلاً من تجميعها، ما عليك سوى تمريرها كخلاصات.

للحصول على تفاصيل عن وحدة ماكرو إصدار tf_library، راجِع tfcompile.bzl.

للحصول على تفاصيل حول أداة tfcompile الأساسية، يُرجى الاطّلاع على tfcompile_main.cc.

الخطوة 3: كتابة التعليمة البرمجية لاستدعاء الرسم البياني الفرعي

تستخدِم هذه الخطوة ملف العنوان (test_graph_tfmatmul.h) الذي تم إنشاؤه من خلال وحدة ماكرو الإصدار tf_library في الخطوة السابقة لاستدعاء الرمز الذي تم إنشاؤه. يقع ملف الرأس في دليل bazel-bin المقابل لحزمة الإنشاء، وتتم تسميته استنادًا إلى سمة الاسم التي تم تعيينها في ماكرو الإنشاء tf_library. على سبيل المثال، العنوان الذي تم إنشاؤه للسمة test_graph_tfmatmul سيكون test_graph_tfmatmul.h. فيما يلي نسخة مختصرة من ما يتم إنشاؤه. يحتوي الملف الذي تم إنشاؤه، في bazel-bin، على تعليقات مفيدة إضافية.

namespace foo {
namespace bar {

// MatMulComp represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code.
class MatMulComp {
 public:
  // AllocMode controls the buffer allocation mode.
  enum class AllocMode {
    ARGS_RESULTS_AND_TEMPS,  // Allocate arg, result and temp buffers
    RESULTS_AND_TEMPS_ONLY,  // Only allocate result and temp buffers
  };

  MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS);
  ~MatMulComp();

  // Runs the computation, with inputs read from arg buffers, and outputs
  // written to result buffers. Returns true on success and false on failure.
  bool Run();

  // Arg methods for managing input buffers. Buffers are in row-major order.
  // There is a set of methods for each positional argument.
  void** args();

  void set_arg0_data(float* data);
  float* arg0_data();
  float& arg0(size_t dim0, size_t dim1);

  void set_arg1_data(float* data);
  float* arg1_data();
  float& arg1(size_t dim0, size_t dim1);

  // Result methods for managing output buffers. Buffers are in row-major order.
  // Must only be called after a successful Run call. There is a set of methods
  // for each positional result.
  void** results();


  float* result0_data();
  float& result0(size_t dim0, size_t dim1);
};

}  // end namespace bar
}  // end namespace foo

تُسمى فئة C++ التي تم إنشاؤها MatMulComp في مساحة الاسم foo::bar، لأنها كانت cpp_class المحدّدة في ماكرو tf_library. تحتوي جميع الفئات التي تم إنشاؤها على واجهة برمجة تطبيقات مشابهة، والفرق الوحيد هو طرق التعامل مع الوسيطة والتخزينات الاحتياطية للنتائج. تختلف هذه الطرق استنادًا إلى عدد المخازن الاحتياطية وأنواعها، التي تم تحديدها من خلال الوسيطتين feed وfetch إلى الماكرو tf_library.

وثمة ثلاثة أنواع من الموارد الاحتياطية المُدارة ضمن الفئة التي تم إنشاؤها: args التي تمثِّل المدخلات وresults تمثِّل المخرجات وtemps التي تمثِّل الموارد الاحتياطية المؤقتة المستخدَمة داخليًا لإجراء العملية الحسابية. وبشكلٍ افتراضي، تخصص كل مثيل للفئة التي تم إنشاؤها كل هذه الموارد الاحتياطية وتديرها نيابة عنك. يمكن استخدام وسيطة الدالة الإنشائية AllocMode لتغيير هذا السلوك. تتم محاذاة جميع المخازن المؤقتة مع حدود 64 بايت.

فئة C++ التي يتم إنشاؤها هي مجرد برنامج تضمين حول التعليمة البرمجية منخفضة المستوى التي يتم إنشاؤها بواسطة XLA.

مثال على استدعاء الدالة التي تم إنشاؤها استنادًا إلى tfcompile_test.cc:

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "third_party/tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());


  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

الخطوة 4: إنشاء البرنامج الثنائي النهائي

تجمع هذه الخطوة بين المكتبة التي تم إنشاؤها بواسطة tf_library في الخطوة 2 والرمز البرمجي المكتوب في الخطوة 3 لإنشاء برنامج ثنائي نهائي. وفي ما يلي مثال على ملف bazel BUILD.

# Example of linking your binary
# Also see //tensorflow/compiler/aot/tests/BUILD
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# The same tf_library call from step 2 above.
tf_library(
    name = "test_graph_tfmatmul",
    ...
)

# The executable code generated by tf_library can then be linked into your code.
cc_binary(
    name = "my_binary",
    srcs = [
        "my_code.cc",  # include test_graph_tfmatmul.h to access the generated header
    ],
    deps = [
        ":test_graph_tfmatmul",  # link in the generated object file
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)