Function ExamplePrepack

tensorflow/lite/experimental/ruy/example_advanced.cc:36–78

};

void ExamplePrepack(ruy::Context* context) {
  const float lhs_data[] = {1, 2, 3, 4};
  const float rhs_data[] = {1, 2, 3, 4};
  float dst_data[4];

  // Set up the matrix layouts and spec.
  ruy::Matrix<float> lhs;
  ruy::MakeSimpleLayout(2, 2, ruy::Order::kRowMajor, &lhs.layout);
  ruy::Matrix<float> rhs;
  ruy::MakeSimpleLayout(2, 2, ruy::Order::kColMajor, &rhs.layout);
  ruy::Matrix<float> dst;
  ruy::MakeSimpleLayout(2, 2, ruy::Order::kColMajor, &dst.layout);
  ruy::BasicSpec<float, float> spec;

  SimpleAllocator allocator;
  auto alloc_fn = [&allocator](std::size_t num_bytes) -> void* {
    return allocator.AllocateBytes(num_bytes);
  };

  // In this example, we pre-pack only the RHS, but either will work.
  // Note that we only need to set the data pointer for the matrix we are
  // pre-packing.
  ruy::PrepackedMatrix prepacked_rhs;
  rhs.data = rhs_data;
  ruy::PrePackForMul<ruy::kAllPaths>(lhs, rhs, spec, context, &dst,
                                     /*prepacked_lhs=*/nullptr, &prepacked_rhs,
                                     alloc_fn);

  // No data will be read from the RHS input matrix when using a pre-packed RHS.
  rhs.data = nullptr;
  lhs.data = lhs_data;
  dst.data = dst_data;
  ruy::MulWithPrepacked<ruy::kAllPaths>(lhs, rhs, spec, context, &dst,
                                        /*prepacked_lhs=*/nullptr,
                                        &prepacked_rhs);
  rhs.data = rhs_data;

  // Print out the results.
  std::cout << "Example Mul with pre-packing RHS, float:\n";
  std::cout << "LHS:\n" << lhs;
  std::cout << "RHS:\n" << rhs;
  std::cout << "Result:\n" << dst << "\n";
}

int main() {
  ruy::Context context;