Add transpose

Author: fujunwei

services/ml/compilation_impl_nn.cc

@@ -33,8 +33,11 @@ CompilationImplNN::CompilationImplNN(const ModelImplNN* model,
 }
 
 CompilationImplNN::~CompilationImplNN() {
-  // ANeuralNetworksCompilation_free(nn_compilation_);
-  // The nn_compilation_ will be deleted in execution phase.
+#if defined(OS_ANDROID)
+  ANeuralNetworksCompilation_free(nn_compilation_);
+#else
+  IE(ie_compilation_free)(ie_compilation_);
+#endif
 }
 
 void CompilationImplNN::Finish(int32_t preference, FinishCallback callback) {

services/ml/execution_impl_nn.cc

@@ -14,23 +14,17 @@
 
 namespace ml {
 
-// TODO:: CompilationImplNN* => std::unique<CompilationImplNN> so that
-// ie_compilation_free(ie_compilation_); can host in class CompilationImplNN.
 ExecutionImplNN::ExecutionImplNN(const CompilationImplNN* compilation,
                                  mojo::ScopedSharedBufferHandle memory)
     : operands_(compilation->operands_),
       operations_(compilation->operations_),
       inputs_(compilation->inputs_),
       outputs_(compilation->outputs_),
       memory_(std::move(memory)),
-#if defined(OS_ANDROID)
-      nn_compilation_(compilation->nn_compilation_) {
-#else
-      ie_compilation_(compilation->ie_compilation_) {
-#endif
+      compilation_impl_(compilation) {
 #if defined(OS_LINUX) || defined(OS_WIN)
   // Create Execution
-  IE(ie_execution_create)(ie_compilation_, &ie_execution_);
+  IE(ie_execution_create)(compilation_impl_->ie_compilation_, &ie_execution_);
 #endif
   uint32_t total_length = 0;
   inputs_info_.reserve(inputs_.size());
@@ -54,9 +48,7 @@ ExecutionImplNN::ExecutionImplNN(const CompilationImplNN* compilation,
 
 ExecutionImplNN::~ExecutionImplNN() {
 #if defined(OS_ANDROID)
-  ANeuralNetworksCompilation_free(nn_compilation_);
 #else
-  IE(ie_compilation_free)(ie_compilation_);
   IE(ie_execution_free)(ie_execution_);
 #endif
   DLOG(INFO) << "ANeuralNetworksCompilation_free";
@@ -91,8 +83,8 @@ void ExecutionImplNN::StartCompute(mojom::UserBufferPtr user_buffer,
   int32_t result = 0;
 #if defined(OS_ANDROID)
   ANeuralNetworksExecution* nn_execution;
-  result =
-      ANeuralNetworksExecution_create(nn_compilation_, &nn_execution);
+  result = ANeuralNetworksExecution_create(compilation_impl_->nn_compilation_,
+                                           &nn_execution);
 #endif
   for (size_t i = 0; i < inputs_info_.size(); ++i) {
     std::unique_ptr<OperandInfo>& info = inputs_info_[i];
@@ -101,8 +93,8 @@ void ExecutionImplNN::StartCompute(mojom::UserBufferPtr user_buffer,
         nn_execution, i, NULL, static_cast<void*>(info->mapping.get()),
         info->length);
 #else
-    result = IE(ie_execution_set_input)(ie_execution_, i,
-                                        info->mapping.get(), info->length);
+    result = IE(ie_execution_set_input)(ie_execution_, i, info->mapping.get(),
+                                        info->length);
 #endif
   }
 
@@ -113,8 +105,8 @@ void ExecutionImplNN::StartCompute(mojom::UserBufferPtr user_buffer,
         nn_execution, i, NULL, static_cast<void*>(info->mapping.get()),
         info->length);
 #else
-    result = IE(ie_execution_set_output)(
-        ie_execution_, i, info->mapping.get(), info->length);
+    result = IE(ie_execution_set_output)(ie_execution_, i, info->mapping.get(),
+                                         info->length);
 #endif
   }
 

services/ml/execution_impl_nn.h

@@ -8,6 +8,7 @@
 #include <vector>
 
 #include "base/macros.h"
+#include "base/memory/scoped_refptr.h"
 #include "services/ml/common.h"
 #include "services/ml/compilation_impl_nn.h"
 #include "services/ml/model_impl_nn.h"
@@ -28,8 +29,7 @@ namespace ml {
 
 class ExecutionImplNN : public mojom::Execution {
  public:
-  ExecutionImplNN(const CompilationImplNN*,
-                       mojo::ScopedSharedBufferHandle);
+  ExecutionImplNN(const CompilationImplNN*, mojo::ScopedSharedBufferHandle);
   ~ExecutionImplNN() override;
 
   void StartCompute(mojom::UserBufferPtr user_buffer,
@@ -46,13 +46,11 @@ class ExecutionImplNN : public mojom::Execution {
   std::vector<std::unique_ptr<OperandInfo>> inputs_info_;
   std::vector<std::unique_ptr<OperandInfo>> outputs_info_;
   mojo::ScopedSharedBufferHandle memory_;
-
+  const CompilationImplNN* compilation_impl_;
 #if defined(OS_LINUX) || defined(OS_WIN)
-  ie_compilation_t* ie_compilation_;
   ie_execution_t* ie_execution_;
-#else
-  ANeuralNetworksCompilation* nn_compilation_;
 #endif
+
   DISALLOW_COPY_AND_ASSIGN(ExecutionImplNN);
 };
 

third_party/blink/renderer/modules/ml/v2/BUILD.gn

@@ -36,6 +36,8 @@ blink_modules_sources("v2") {
     "ops/reshape.h",
     "ops/softmax.cc",
     "ops/softmax.h",
+    "ops/transpose.cc",
+    "ops/transpose.h",
   ]
 
   public_deps = [

third_party/blink/renderer/modules/ml/v2/nn_compilation.cpp

@@ -70,7 +70,7 @@ void NNCompilation::OnCreateExecution(
 
   if (result_code == ml::mojom::blink::NOT_ERROR) {
     resolver->Resolve(MakeGarbageCollected<NNExecution>(
-        std::move(init_params), std::move(name_index_)));
+        std::move(init_params), name_index_));
   } else {
     resolver->Reject(MakeGarbageCollected<DOMException>(
         DOMExceptionCode::kInvalidStateError,

third_party/blink/renderer/modules/ml/v2/nn_context.cpp

@@ -23,6 +23,7 @@
 #include "third_party/blink/renderer/modules/ml/v2/ops/relu.h"
 #include "third_party/blink/renderer/modules/ml/v2/ops/reshape.h"
 #include "third_party/blink/renderer/modules/ml/v2/ops/softmax.h"
+#include "third_party/blink/renderer/modules/ml/v2/ops/transpose.h"
 #include "third_party/blink/renderer/platform/bindings/exception_code.h"
 
 namespace blink {
@@ -237,6 +238,10 @@ Operand* NNContext::matmul(Operand* a, Operand* b) {
   return MakeGarbageCollected<MatMul>(a, b);
 }
 
+Operand* NNContext::transpose(Operand* input, WTF::Vector<int32_t> new_shape) {
+  return MakeGarbageCollected<Transpose>(input, std::move(new_shape));
+}
+
 ScriptPromise NNContext::createModel(ScriptState* script_state,
                                      const NamedOperandVector& outputs) {
   auto* resolver = MakeGarbageCollected<ScriptPromiseResolver>(script_state);

third_party/blink/renderer/modules/ml/v2/nn_context.h

@@ -71,6 +71,7 @@ class NNContext final : public ScriptWrappable,
   Operand* relu(Operand*);
   Operand* matmul(Operand*, Operand*);
   ScriptPromise createModel(ScriptState*, const NamedOperandVector&);
+  Operand* transpose(Operand*, WTF::Vector<int32_t>);
 
   // ExecutionContextLifecycleObserver overrides.
   void ContextDestroyed() override;

third_party/blink/renderer/modules/ml/v2/nn_context.idl

@@ -28,6 +28,7 @@ interface NNContext {
     Operand softmax(Operand input);
     Operand relu(Operand input);
     Operand matmul(Operand a, Operand b);
+    Operand transpose(Operand input, optional sequence<long> permutation=[]);
 
     // Create Model
     [CallWith=ScriptState] Promise<Model> createModel(sequence<NamedOperand> outputs);

third_party/blink/renderer/modules/ml/v2/ops/transpose.cc

@@ -0,0 +1,41 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "third_party/blink/renderer/modules/ml/v2/ops/transpose.h"
+
+#include <memory>
+
+#include "third_party/blink/renderer/modules/ml/neural_network_context.h"
+
+namespace blink {
+
+Transpose::Transpose(Operand* input, WTF::Vector<int32_t> permutation)
+    : Output({input}), permutation_(permutation) {}
+
+void Transpose::AddLayer(NNModel* model, uint32_t& index) {
+  Vector<uint32_t> input_indexes;
+  // Add input index to input_indexes.
+  for (auto& input : Output::Inputs()) {
+    input_indexes.push_back(input->Index());
+  }
+
+  // Add permutation operand and set the value.
+  if (!permutation_.IsEmpty()) {
+    uint32_t permutation_index = index++;
+    // The new shape is 1-D tensor.
+    Vector<uint32_t> permutation_dims(1, permutation_.size());
+    model->AddTensorOperand(permutation_index, permutation_dims, permutation_);
+    input_indexes.push_back(permutation_index);
+  }
+
+  // Add Reshape output operand.
+  uint32_t output_index = index++;
+  Operand::SetIndex(output_index);
+  model->AddUnspecifiedOperand();
+
+  model->AddOperation(NeuralNetworkContext::kTranspose, input_indexes,
+                      {output_index});
+}
+
+}  // namespace blink

third_party/blink/renderer/modules/ml/v2/ops/transpose.h

@@ -0,0 +1,28 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef THIRD_PARTY_BLINK_RENDERER_MODULES_ML_OPS_TRANSPOSE_H_
+#define THIRD_PARTY_BLINK_RENDERER_MODULES_ML_OPS_TRANSPOSE_H_
+
+#include "third_party/blink/renderer/modules/ml/v2/nn_model.h"
+#include "third_party/blink/renderer/modules/ml/v2/operand.h"
+#include "third_party/blink/renderer/modules/ml/v2/ops/output.h"
+#include "third_party/blink/renderer/platform/wtf/vector.h"
+
+namespace blink {
+
+class Transpose final : public Output {
+ public:
+  Transpose(Operand*, WTF::Vector<int32_t>);
+  ~Transpose() override = default;
+
+  void AddLayer(NNModel* model, uint32_t& index) override;
+
+ private:
+  Vector<int32_t> permutation_;
+};
+
+}  // namespace blink
+
+#endif  // THIRD_PARTY_BLINK_RENDERER_MODULES_ML_OPS_TRANSPOSE_H_