onehot vectorized (#7339)

This commit improves performance of one_hot custom op

libnd4j/include/ops/declarable/generic/parity_ops/onehot.cpp

@@ -23,6 +23,7 @@
 
 #include <ops/declarable/CustomOperations.h>
 #include <helpers/ShapeUtils.h>
+#include <ops/declarable/helpers/one_hot.h>
 
 namespace nd4j {
     namespace ops {
@@ -65,20 +66,8 @@ namespace nd4j {
             if (axis < 0)
                 axis = output->rankOf() + axis;
 
-            auto vec = ShapeUtils::convertAxisToTadTarget(input->rankOf(), {axis});
-            auto tads = output->allTensorsAlongDimension({axis});
-            for (int e = 0; e < tads->size(); e++) {
-                auto tad = tads->at(e);
-                tad->assign(off);
-
-                int idx = input->e<int>(e);
-                if (idx < 0 || idx >= tad->lengthOf())
-                    continue;
-
-                tad->p(idx, on);
-            }
-
-            delete tads;
+            std::vector<int> vec({axis});
+            helpers::onehot(output, input, vec, on, off);
 
             return Status::OK();
         }

libnd4j/include/ops/declarable/helpers/cpu/one_hot.cpp

@@ -0,0 +1,78 @@
+/*******************************************************************************
+ * Copyright (c) 2015-2018 Skymind, Inc.
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+
+#include <helpers/TAD.h>
+#include "../one_hot.h"
+
+namespace nd4j {
+    namespace ops {
+        namespace helpers {
+            template <typename Z, typename I>
+            static void onehot_(void *voutput, Nd4jLong *zShapeInfo, void *vindices, Nd4jLong *iShapeInfo, std::vector<int> &axis, double on, double off) {
+                auto output = reinterpret_cast<Z*>(voutput);
+                auto indices = reinterpret_cast<I*>(vindices);
+
+                shape::TAD tad;
+                tad.init(zShapeInfo, axis.data(), axis.size());
+                tad.createTadOnlyShapeInfo();
+                tad.createOffsets();
+                auto iLen = static_cast<unsigned int>(shape::length(iShapeInfo));
+                auto tLen = static_cast<unsigned int>(shape::length(tad.tadOnlyShapeInfo));
+                auto numTads = static_cast<unsigned int>(tad.numTads);
+
+                nd4j_printf("numTads: [%i]; iLen: [%i]\n", numTads, iLen);
+
+                if (iLen != numTads)
+                    throw std::runtime_error("OneHot: number of TADs should be equal to number of indices");
+
+                if (shape::elementWiseStride(zShapeInfo) != 1 || shape::elementWiseStride(zShapeInfo) != 1)
+                    throw std::runtime_error("OneHot: op expects output and indices to have elementWiseStride to be equal to 1");
+
+                Z zero = static_cast<Z>(off);
+                Z one = static_cast<Z>(on);
+
+                PRAGMA_OMP_PARALLEL_FOR
+                for (unsigned int e = 0; e < numTads; e++) {
+                    auto cO = output + tad.tadOffsets[e];
+
+                    auto idx = static_cast<unsigned int>(indices[e]);
+                    if (idx < 0 || idx >= tLen) {
+                        PRAGMA_OMP_SIMD
+                        for (unsigned int t = 0; t < tLen; t++) {
+                            cO[t] = zero;
+                        }
+                    } else {
+                        PRAGMA_OMP_SIMD
+                        for (unsigned int t = 0; t < tLen; t++) {
+                            cO[t] = idx == t ? one : zero;
+                        }
+                    }
+                }
+            }
+
+            void onehot(NDArray *output, NDArray *indices, std::vector<int> &axis, double on, double off) {
+                auto zType = output->dataType();
+                auto iType = indices->dataType();
+
+                BUILD_DOUBLE_SELECTOR(zType, iType, onehot_, (output->buffer(), output->shapeInfo(), indices->buffer(), indices->shapeInfo(), axis, on, off), LIBND4J_TYPES, LIBND4J_TYPES);
+            }
+        }
+    }
+}

libnd4j/include/ops/declarable/helpers/one_hot.h

@@ -0,0 +1,35 @@
+/*******************************************************************************
+ * Copyright (c) 2015-2018 Skymind, Inc.
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+
+#ifndef DEV_TESTS_ONE_HOT_H
+#define DEV_TESTS_ONE_HOT_H
+
+#include <op_boilerplate.h>
+#include <NDArray.h>
+
+namespace nd4j {
+    namespace ops {
+        namespace helpers {
+            void onehot(NDArray *output, NDArray *indices, std::vector<int> &axis, double on, double off);
+        }
+    }
+}
+
+#endif //DEV_TESTS_ONE_HOT_H