MurmurHash3 operator (#174)

* MurmurHash3 operator

onnxruntime/contrib_ops/contrib_kernels.cc

@@ -17,6 +17,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3);
 
 void RegisterContribKernels(std::function<void(KernelCreateInfo&&)> fn) {
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp)>());
@@ -32,6 +33,7 @@ void RegisterContribKernels(std::function<void(KernelCreateInfo&&)> fn) {
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer)>());
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression)>());
   fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range)>());
+  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3)>());
 }
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/murmur_hash3.cc

@@ -0,0 +1,202 @@
+//-----------------------------------------------------------------------------
+// MurmurHash3 was written by Austin Appleby, and is placed in the public
+// domain. The author hereby disclaims copyright to this source code.
+
+// License from https://github.com/scikit-learn/scikit-learn/blob/master/COPYING
+/* Modifications Copyright (c) Microsoft. */
+
+#include "contrib_ops/cpu/murmur_hash3.h"
+
+// Platform-specific functions and macros
+
+// Microsoft Visual Studio
+
+#if defined(_MSC_VER)
+
+#define FORCE_INLINE __forceinline
+
+#include <stdlib.h>
+
+#define ROTL32(x, y) _rotl(x, y)
+#define ROTL64(x, y) _rotl64(x, y)
+
+#define BIG_CONSTANT(x) (x)
+
+// Other compilers
+
+#else  // defined(_MSC_VER)
+
+#if defined(GNUC) && ((GNUC > 4) || (GNUC == 4 && GNUC_MINOR >= 4))
+
+// gcc version >= 4.4 4.1 = RHEL 5, 4.4 = RHEL 6.
+// Don't inline for RHEL 5 gcc which is 4.1
+#define FORCE_INLINE attribute((always_inline))
+
+#else
+
+#define FORCE_INLINE
+
+#endif
+
+inline uint32_t rotl32(uint32_t x, int8_t r) {
+  return (x << r) | (x >> (32 - r));
+}
+
+inline uint64_t rotl64(uint64_t x, int8_t r) {
+  return (x << r) | (x >> (64 - r));
+}
+
+#define ROTL32(x, y) rotl32(x, y)
+#define ROTL64(x, y) rotl64(x, y)
+
+#define BIG_CONSTANT(x) (x##LLU)
+
+#endif  // !defined(_MSC_VER)
+
+//-----------------------------------------------------------------------------
+// If your platform needs to do endian-swapping or can only
+// handle aligned reads, do the conversion here
+
+FORCE_INLINE uint32_t getblock(const uint32_t* p, int i) {
+  return p[i];
+}
+
+FORCE_INLINE uint64_t getblock(const uint64_t* p, int i) {
+  return p[i];
+}
+
+//-----------------------------------------------------------------------------
+// Finalization mix - force all bits of a hash block to avalanche
+
+FORCE_INLINE uint32_t fmix(uint32_t h) {
+  h ^= h >> 16;
+  h *= 0x85ebca6b;
+  h ^= h >> 13;
+  h *= 0xc2b2ae35;
+  h ^= h >> 16;
+
+  return h;
+}
+
+//----------
+
+FORCE_INLINE uint64_t fmix(uint64_t k) {
+  k ^= k >> 33;
+  k *= BIG_CONSTANT(0xff51afd7ed558ccd);
+  k ^= k >> 33;
+  k *= BIG_CONSTANT(0xc4ceb9fe1a85ec53);
+  k ^= k >> 33;
+
+  return k;
+}
+
+
+namespace onnxruntime {
+namespace contrib {
+
+ONNX_OPERATOR_KERNEL_EX(
+    MurmurHash3,
+    kMSDomain,
+    1,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .TypeConstraint("T1", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
+                                                      DataTypeImpl::GetTensorType<uint32_t>(),
+                                                      DataTypeImpl::GetTensorType<std::string>()})
+        .TypeConstraint("T2", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
+                                                      DataTypeImpl::GetTensorType<uint32_t>()}),
+    MurmurHash3);
+
+void MurmurHash3::MurmurHash3_x86_32(const void* key, int len, uint32_t seed, void* out) const {
+
+  const uint8_t* data = reinterpret_cast<const uint8_t*>(key);
+  const int nblocks = len / 4;
+  uint32_t h1 = seed;
+  const uint32_t c1 = 0xcc9e2d51;
+  const uint32_t c2 = 0x1b873593;
+
+  //----------
+  // body
+  const uint32_t* blocks = reinterpret_cast<const uint32_t*>(data + static_cast<int64_t>(nblocks) * 4);
+
+  for (int i = -nblocks; i; i++) {
+    uint32_t k1 = getblock(blocks, i);
+
+    k1 *= c1;
+    k1 = ROTL32(k1, 15);
+    k1 *= c2;
+
+    h1 ^= k1;
+    h1 = ROTL32(h1, 13);
+    h1 = h1 * 5 + 0xe6546b64;
+  }
+
+  //----------
+  // tail
+  const uint8_t* tail = reinterpret_cast<const uint8_t*>(data + static_cast<int64_t>(nblocks) * 4);
+
+  uint32_t k1 = 0;
+
+  switch (len & 3) {
+    case 3:
+      k1 ^= tail[2] << 16;
+    case 2:
+      k1 ^= tail[1] << 8;
+    case 1:
+      k1 ^= tail[0];
+      k1 *= c1;
+      k1 = ROTL32(k1, 15);
+      k1 *= c2;
+      h1 ^= k1;
+  };
+
+  //----------
+  // finalization
+  h1 ^= len;
+
+  h1 = fmix(h1);
+
+  *(uint32_t*)out = h1;
+}
+
+Status MurmurHash3::Compute(OpKernelContext* ctx) const {
+  const Tensor* keys = ctx->Input<Tensor>(0);
+  ONNXRUNTIME_ENFORCE(keys);
+
+  const TensorShape& input_shape = keys->Shape();
+  Tensor* output_tensor = ctx->Output(0, input_shape);
+
+  const MLDataType keys_type = keys->DataType();
+  const int input_element_bytes = static_cast<int>(keys->DataType()->Size());
+  const int output_element_bytes = static_cast<int>(output_tensor->DataType()->Size());
+  const int64_t input_count = input_shape.Size();
+  for (int i = 0; i < input_count; ++i) {
+    if (DataTypeImpl::GetType<std::string>() == keys_type) {
+      auto input = keys->DataRaw();
+      auto output = output_tensor->MutableDataRaw();
+      auto input_string = reinterpret_cast<const std::string*>(input)[i];
+      MurmurHash3_x86_32(input_string.c_str(),
+                         static_cast<int>(input_string.length()),
+                         seed_,
+                         reinterpret_cast<uint32_t*>(output) + static_cast<int64_t>(i) * output_element_bytes);
+    } else {
+      auto output_type = output_tensor->DataType();
+      if ((DataTypeImpl::GetType<int32_t>() == keys_type || DataTypeImpl::GetType<uint32_t>() == keys_type) &&
+          (DataTypeImpl::GetType<int32_t>() == output_type || DataTypeImpl::GetType<uint32_t>() == output_type)) {
+        auto input = keys->DataRaw();
+        auto output = output_tensor->MutableDataRaw();
+        MurmurHash3_x86_32(reinterpret_cast<const uint8_t*>(input) + static_cast<int64_t>(i) * input_element_bytes,
+                           input_element_bytes,
+                           seed_,
+                           reinterpret_cast<uint8_t*>(output) + static_cast<int64_t>(i) * output_element_bytes);
+      } else {
+        return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "Type not supported.");
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/murmur_hash3.h

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+class MurmurHash3 final : public OpKernel {
+ public:
+  MurmurHash3(const OpKernelInfo& info) : OpKernel(info) {
+    seed_ = static_cast<uint32_t>(info.GetAttrOrDefault<int64_t>("seed", 0));
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+private:
+  void MurmurHash3_x86_32(const void* key, int len, uint32_t seed, void* out) const;
+
+private :
+  uint32_t seed_;
+};
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -528,6 +528,28 @@ The bounding box coordinates corresponding to the selected indices can then be o
         }
       });
 
+  ONNX_CONTRIB_OPERATOR_SCHEMA(MurmurHash3)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(R"DOC(The underlying implementation is MurmurHash3_x86_32 generating low latency 32bits hash suitable for implementing lookup tables, Bloom filters, count min sketch or feature hashing.)DOC")
+      .Input(0, "X", "An input tensor to hash.", "T1")
+      .Output(0, "Y", "32-bit hash value.", "T2")
+      .TypeConstraint("T1", {"tensor(uint32)", "tensor(int32)", "tensor(string)"}, "Constrain input type to unsigned or signed 32-bit integer tensor, or string tensor. It should be utf-8 encoded if using unicode.")
+      .TypeConstraint("T2", {"tensor(uint32)", "tensor(int32)"}, "Constrain output type to unsigned or signed 32-bit integer tensor.")
+      .Attr(
+          "seed",
+          "Seed for the hashing algorithm, unsigned 32-bit integer, default to 0.",
+          AttributeProto::INT,
+          (int64_t)0LL)
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        // Shape inference
+        if (!hasInputShape(ctx, 0))
+          return;
+
+        auto& input_shape = getInputShape(ctx, 0);
+        updateOutputShape(ctx, 0, input_shape);
+      });
+
   ONNX_CONTRIB_OPERATOR_SCHEMA(StringNormalizer)
       .SetDomain(kMSDomain)
       .SinceVersion(1)

onnxruntime/test/contrib_ops/murmur_hash3_test.cc

@@ -0,0 +1,98 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(MurmurHash3OpTest, DefaultSeed) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {3L});
+  test.AddOutput<int32_t>("Y", {1}, {847579505L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, ZeroSeed) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {3L});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<int32_t>("Y", {1}, {847579505L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, ZeroSeedUIntResult) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {3L});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<uint32_t>("Y", {1}, {847579505L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, ZeroSeedUIntResult2) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {4L});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<uint32_t>("Y", {1}, {1889779975L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, MoreData) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {2}, {3L, 4L});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<uint32_t>("Y", {2}, {847579505L, 1889779975L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest,NonZeroSeed) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {3L});
+  test.AddAttribute<int64_t>("seed", 42LL);
+  test.AddOutput<int32_t>("Y", {1}, {-1823081949L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, NonZeroSeedUIntResult) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<int32_t>("X", {1}, {3L});
+  test.AddAttribute<int64_t>("seed", 42LL);
+  test.AddOutput<uint32_t>("Y", {1}, {2471885347L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, StringKeyIntResult) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<std::string>("X", {1}, {"foo"});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<int32_t>("Y", {1}, {-156908512L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, StringKeyUIntResult) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<std::string>("X", {1}, {"foo"});
+  test.AddAttribute<int64_t>("seed", 0LL);
+  test.AddOutput<uint32_t>("Y", {1}, {4138058784L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, StringKeyIntWithSeed42) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<std::string>("X", {1}, {"foo"});
+  test.AddAttribute<int64_t>("seed", 42LL);
+  test.AddOutput<int32_t>("Y", {1}, {-1322301282L});
+  test.Run();
+}
+
+TEST(MurmurHash3OpTest, StringKeyUIntWithSeed42) {
+  OpTester test("MurmurHash3", 1, onnxruntime::kMSDomain);
+  test.AddInput<std::string>("X", {1}, {"foo"});
+  test.AddAttribute<int64_t>("seed", 42LL);
+  test.AddOutput<uint32_t>("Y", {1}, {2972666014L});
+  test.Run();
+}
+
+}  // namespace test
+}  // namespace onnxruntime