External Operators 2

CMakeLists.txt

@@ -89,6 +89,7 @@ option(USE_TENSORRT "Enable inference optimization with TensorRT." OFF)
 option(USE_ASAN "Enable Clang/GCC ASAN sanitizers." OFF)
 cmake_dependent_option(ENABLE_TESTCOVERAGE "Enable compilation with test coverage metric output" OFF "NOT MSVC" OFF)
 option(USE_INT64_TENSOR_SIZE "Use int64_t to represent the total number of elements in a tensor" OFF)
+option(BUILD_EXTENSION "Path to extension to build" "")
 option(BUILD_CYTHON_MODULES "Build cython modules." OFF)
 option(LOG_FATAL_THROW "Log exceptions but do not abort" ON)
 cmake_dependent_option(USE_SPLIT_ARCH_DLL "Build a separate DLL for each Cuda arch (Windows only)." ON "MSVC" OFF)
@@ -787,6 +788,10 @@ add_library(transposerowsp_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/example/extens
 add_library(subgraph_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/example/extensions/lib_subgraph/subgraph_lib.cc ${CMAKE_CURRENT_SOURCE_DIR}/src/lib_api.cc)
 add_library(pass_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/example/extensions/lib_pass/pass_lib.cc ${CMAKE_CURRENT_SOURCE_DIR}/src/lib_api.cc)
 
+if(IS_DIRECTORY ${BUILD_EXTENSION})
+  add_subdirectory(${BUILD_EXTENSION} ${BUILD_EXTENSION}/build)
+endif()
+
 target_include_directories(customop_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include/mxnet)
 target_include_directories(transposecsr_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include/mxnet)
 target_include_directories(transposerowsp_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include/mxnet)

example/extensions/lib_external_ops/CMakeLists.txt

@@ -0,0 +1,27 @@
+# specify CXX sources
+FILE(GLOB CXX_SRCS
+  # Required files
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../../src/lib_api.cc
+  # Your custom files
+  ${CMAKE_CURRENT_SOURCE_DIR}/init_lib.cc
+  ${CMAKE_CURRENT_SOURCE_DIR}/min_ex.cc
+  )
+
+# create library & set libraries
+add_library(external_lib SHARED ${CXX_SRCS})
+target_link_libraries(external_lib PUBLIC mxnet)
+
+# generic GPU stuff
+if(USE_CUDA)
+  # specify GPU sources (optional)
+  FILE(GLOB CU_SRCS "*.cu")
+  target_sources(external_lib PUBLIC ${CU_SRCS})
+endif(USE_CUDA)
+
+if(UNIX)
+  # unix-specific stuff
+  if(USE_CUDA)
+    # unix+GPU-specific stuff
+    target_compile_options(external_lib PUBLIC -shared)
+  endif(USE_CUDA)
+endif()

example/extensions/lib_external_ops/README.md

@@ -0,0 +1,70 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License.  You may obtain a copy of the License at -->
+
+<!---   http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied.  See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+
+External Operators Example and Tutorial
+=======================================
+
+## Introduction
+
+Extending MXNet with custom components used to mean distributing a custom fork. This feature allows adding custom components to MXNet by dynamically loading external libraries at runtime. Currently it is only supported on Linux systems (Windows and Mac are __NOT__ supported). 
+
+## Getting Started
+
+### Have MXNet Ready
+
+For this tutorial, clone MXNet from source like:
+```
+git clone https://github.com/apache/incubator-mxnet.git --recursive --init
+```
+
+Build MXNet like:
+```
+cp config/linux.cmake config.cmake
+mkdir build
+cd build
+cmake ..
+cmake --build .
+```
+
+## Run An Example
+
+This example shows compiling a custom backend operator and then dynamically loading it into MXNet at runtime. Go to the **lib_external_ops** directory and follow these steps:
+
+1. Touch or modify the **min_ex.cc** and/or **min_ex-inl.h** file(s)
+2. Go into the **build** directory that was created when building MXNet.
+3. Run `cmake .. -DBUILD_EXTENSION=$(pwd)/../example/extensions/lib_external_ops`
+4. Run `cmake --build .`
+5. Go to the **example/extensions/lib_external_ops** directory again
+6. Run `python test_loading.py` to execute the test program. You should see the following output:
+```
+Operator not registered yet
+MXNet version 20000 supported
+[]
+Operator executed successfully
+```
+
+## Writing an External Operator Library
+To build your own library containing custom components, compose a C++ source file like `mycomp_lib.cc`, include the `lib_api.h` header file, compile the `lib_api.cc` file, and implement the following required function:
+- `initialize` - Library Initialization Function
+
+Then create a CMakeLists.txt file and set `mxnet` as a link library like:
+```
+add_library(external_lib SHARED ${SRCS})
+target_link_libraries(external_lib PUBLIC mxnet)
+```
+
+Next, build MXNet and set the path to your directory with the CMakeLists.txt file via the `BUILD_EXTENSION` option. This will build your library with all of the MXNet includes. 

example/extensions/lib_external_ops/init_lib.cc

@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2020 by Contributors
+ * \file init_lib.cc
+ * \brief initialize function implementation library file
+ */
+
+#include <iostream>
+#include "mxnet/lib_api.h"
+
+using namespace mxnet::ext;
+
+MXReturnValue initialize(int version) {
+  if (version >= 10700) {
+    std::cout << "MXNet version " << version << " supported" << std::endl;
+    return MX_SUCCESS;
+  } else {
+    MX_ERROR_MSG << "MXNet version " << version << " not supported";
+    return MX_FAIL;
+  }
+}

example/extensions/lib_external_ops/min_ex-inl.h

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2020 by Contributors
+ * \file min_ex-inl.h
+ * \brief example external operator header file
+ */
+
+#ifndef MXNET_OPERATOR_TENSOR_MIN_EX_OP_INL_H_
+#define MXNET_OPERATOR_TENSOR_MIN_EX_OP_INL_H_
+
+#include <dmlc/parameter.h>
+#include <vector>
+#include <algorithm>
+#include "operator/mxnet_op.h"
+#include "operator/operator_common.h"
+#include "operator/elemwise_op_common.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename xpu>
+void MinExForward(const nnvm::NodeAttrs& attrs,
+                  const OpContext& ctx,
+                  const std::vector<TBlob>& inputs,
+                  const std::vector<OpReqType>& req,
+                  const std::vector<TBlob>& outputs) {
+  //do nothing                                                                                                                                                                         
+}
+
+
+inline bool MinExOpShape(const nnvm::NodeAttrs& attrs,
+                         mxnet::ShapeVector* in_attrs,
+                         mxnet::ShapeVector* out_attrs) {
+    //do nothing                                                                                                                                                                       
+    return true;
+}
+
+inline bool MinExOpType(const nnvm::NodeAttrs& attrs,
+                        std::vector<int> *in_attrs,
+                        std::vector<int> *out_attrs) {
+  //do nothing                                                                                                                                                                         
+  return true;
+}
+
+}  // namespace op                                                                                                                                                                     
+}  // namespace mxnet                                                                                                                                                                  
+
+#endif  // MXNET_OPERATOR_TENSOR_MIN_EX_OP_INL_H_

example/extensions/lib_external_ops/min_ex.cc

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2020 by Contributors
+ * \file min_ex.cc
+ * \brief example external operator source file
+ */
+
+#include "min_ex-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(min_ex)
+.describe("some description")
+.set_num_inputs(0)
+.set_num_outputs(0)
+.set_attr<mxnet::FInferShape>("FInferShape", MinExOpShape)
+.set_attr<nnvm::FInferType>("FInferType", MinExOpType)
+.set_attr<FCompute>("FCompute<cpu>", MinExForward<cpu>);
+
+}  // namespace op                                                                                                                                                                     
+}  // namespace mxnet 

example/extensions/lib_external_ops/min_ex.cu

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2020 by Contributors
+ * \file min_ex.cu
+ * \brief example external operator CUDA source file
+ */
+
+#include "./min_ex-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(min_ex)
+.set_attr<FCompute>("FCompute<gpu>", MinExForward<gpu>);
+
+}  // namespace op
+}  // namespace mxnet

example/extensions/lib_external_ops/test_loading.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# coding: utf-8
+# pylint: disable=arguments-differ
+
+# This test checks if dynamic loading of library into MXNet is successful
+# and checks the computation of an external operator
+
+import mxnet as mx
+import os
+
+# check if operator exists
+if hasattr(mx.nd,'min_ex'):
+    raise Exception('Operator already loaded')
+else:
+    print('Operator not registered yet')
+
+# test loading library
+if (os.name=='posix'):
+    path = os.path.abspath('build/libexternal_lib.so')
+    mx.library.load(path,False)
+
+# execute operator
+print(mx.nd.min_ex())
+print('Operator executed successfully')