Add vm.buffer.hash and util.buffer.hash ops (#16003)

This adds support for computing the SipHash24 of a byte range of a buffer, with the goal of using this as a primitive for building a hash table out of VM IR.
iree-org · Jan 9, 2024 · a7a7ad6 · a7a7ad6
1 parent 4b2aaaf
commit a7a7ad6
Show file tree

Hide file tree

Showing 16 changed files with 324 additions and 4 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.cpp
@@ -1506,6 +1506,25 @@ void BufferStoreOp::setSubrangeOperand(unsigned operandIndex,
   getLengthMutable().assign(operand.length);
 }
 
+SubrangeOperand BufferHashOp::getSubrangeOperand(unsigned operandIndex) {
+  if (operandIndex == 0) {
+    return SubrangeOperand{getSource(), getSourceSize(), getSourceOffset(),
+                           getLength()};
+  } else {
+    assert(false && "only source is a subrange");
+    return {};
+  }
+}
+
+void BufferHashOp::setSubrangeOperand(unsigned operandIndex,
+                                      SubrangeOperand operand) {
+  assert(operandIndex == 0 && "only source is a subrange");
+  getSourceMutable().assign(operand.resource);
+  getSourceSizeMutable().assign(operand.resourceSize);
+  getSourceOffsetMutable().assign(operand.offset);
+  getLengthMutable().assign(operand.length);
+}
+
 } // namespace mlir::iree_compiler::IREE::Util
 
 #define GET_OP_CLASSES

diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.td b/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.td
@@ -1341,6 +1341,40 @@ def Util_BufferStoreOp : Util_Op<"buffer.store", [
   }];
 }
 
+def Util_BufferHashOp : Util_Op<"buffer.hash", [
+  MemoryEffects<[MemRead]>,
+  Util_SizeAwareOp,
+  DeclareOpInterfaceMethods<Util_SubrangeOperandOpInterface>,
+]> {
+  let summary = [{computes the hash of a byte range of a buffer}];
+  let description = [{
+    Computes the SipHash-2-4 of a value at a byte offset with the given length.
+    This always uses a seed of `0x0001020304...0e0f` and produces a single 64
+    bit value.
+  }];
+
+  let arguments = (ins
+    Util_BufferType:$source,
+    Util_Size:$source_size,
+    Util_Offset:$source_offset,
+    Util_Size:$length
+  );
+  let results = (outs
+    I64:$result
+  );
+
+  let assemblyFormat = [{
+    $source `[` $source_offset `for` $length `]`
+    `:` type($source) `` `{` $source_size `}` `->` type($result)
+    attr-dict-with-keyword
+  }];
+
+  let extraClassDeclaration = [{
+    Value getOperandSize(unsigned idx) { return getSourceSize(); }
+    Value getResultSize(unsigned idx) { return {}; }
+  }];
+}
+
 } // OpGroupBufferOps
 
 //===----------------------------------------------------------------------===//

diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/test/buffer_ops.mlir b/compiler/src/iree/compiler/Dialect/Util/IR/test/buffer_ops.mlir
@@ -126,3 +126,14 @@ func.func @buffer_store(%arg0: !util.buffer, %arg1: index, %arg2: i32) {
   util.buffer.store %arg2, %arg0[%c100 for %c4] : i32 -> !util.buffer{%arg1}
   return
 }
+
+// -----
+
+// CHECK-LABEL: @buffer_hash
+func.func @buffer_hash(%arg0: !util.buffer, %arg1: index) -> i64 {
+  %c17 = arith.constant 17 : index
+  %c100 = arith.constant 100 : index
+  // CHECK: = util.buffer.hash %arg0[%c100 for %c17] : !util.buffer{%arg1} -> i64
+  %0 = util.buffer.hash %arg0[%c100 for %c17] : !util.buffer{%arg1} -> i64
+  return %0 : i64
+}
diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/UtilToVM/ConvertBufferOps.cpp b/compiler/src/iree/compiler/Dialect/VM/Conversion/UtilToVM/ConvertBufferOps.cpp
@@ -318,6 +318,22 @@ struct BufferStoreOpConversion
   }
 };
 
+struct BufferHashOpConversion
+    : public OpConversionPattern<IREE::Util::BufferHashOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(IREE::Util::BufferHashOp hashOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto newType =
+        getTypeConverter()->convertType(hashOp.getResult().getType());
+    auto byteOffset = castToI64(adaptor.getSourceOffset(), rewriter);
+    auto length = castToI64(adaptor.getLength(), rewriter);
+    rewriter.replaceOpWithNewOp<IREE::VM::BufferHashOp>(
+        hashOp, newType, adaptor.getSource(), byteOffset, length);
+    return success();
+  }
+};
+
 } // namespace
 
 void populateUtilBufferToVMPatterns(MLIRContext *context,
@@ -341,7 +357,8 @@ void populateUtilBufferToVMPatterns(MLIRContext *context,
                     IREE::Util::BufferDeallocOp, IREE::Util::BufferSliceOp,
                     IREE::Util::BufferSizeOp, IREE::Util::BufferCopyOp,
                     IREE::Util::BufferCompareOp, IREE::Util::BufferFillOp,
-                    IREE::Util::BufferLoadOp, IREE::Util::BufferStoreOp>();
+                    IREE::Util::BufferLoadOp, IREE::Util::BufferStoreOp,
+                    IREE::Util::BufferHashOp>();
 
   patterns.insert<BufferConstantOpConversion>(typeConverter, context);
   patterns.insert<BufferAllocOpConversion>(typeConverter, context);
@@ -353,6 +370,7 @@ void populateUtilBufferToVMPatterns(MLIRContext *context,
   patterns.insert<BufferFillOpConversion>(typeConverter, context);
   patterns.insert<BufferLoadOpConversion>(typeConverter, context);
   patterns.insert<BufferStoreOpConversion>(typeConverter, context);
+  patterns.insert<BufferHashOpConversion>(typeConverter, context);
 }
 
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/buffer_ops.mlir b/compiler/src/iree/compiler/Dialect/VM/Conversion/UtilToVM/test/buffer_ops.mlir
@@ -295,3 +295,20 @@ func.func @buffer_store_index(%arg0: !util.buffer, %arg1: index, %arg2: index) {
   util.buffer.store %arg2, %arg0[%byte_offset for %element_size] : index -> !util.buffer{%arg1}
   return
 }
+
+// -----
+
+// CHECK-LABEL: @buffer_hash
+func.func @buffer_hash(%arg0: !util.buffer, %arg1: index) -> i64 {
+  %byte_offset = arith.constant 128 : index
+  %length = arith.constant 17 : index
+  // CHECK-32-DAG: %[[BYTE_OFFSET:.+]] = vm.const.i64 128
+  // CHECK-32-DAG: %[[LENGTH:.+]] = vm.const.i64 17
+  // CHECK-32: %[[VALUE:.+]] = vm.buffer.hash %arg0, %[[BYTE_OFFSET]], %[[LENGTH]] : !vm.buffer -> i64
+  // CHECK-64-DAG: %[[BYTE_OFFSET:.+]] = vm.const.i64 128
+  // CHECK-64-DAG: %[[LENGTH:.+]] = vm.const.i64 17
+  // CHECK-64: %[[VALUE:.+]] = vm.buffer.hash %arg0, %[[BYTE_OFFSET]], %[[LENGTH]] : !vm.buffer -> i64
+  %0 = util.buffer.hash %arg0[%byte_offset for %length] : !util.buffer{%arg1} -> i64
+  // CHECK: return %[[VALUE]]
+  return %0 : i64
+}
diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.cpp b/compiler/src/iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.cpp
@@ -4518,6 +4518,8 @@ void populateVMToEmitCPatterns(ConversionTarget &conversionTarget,
                         DenseSet<size_t>({0}), true);
   ADD_CONTAINER_PATTERN(IREE::VM::BufferStoreI64Op, "vm_buffer_store_i64",
                         DenseSet<size_t>({0}), true);
+  ADD_CONTAINER_PATTERN(IREE::VM::BufferHashOp, "iree_vm_buffer_hash",
+                        DenseSet<size_t>({0}), true);
   ADD_CONTAINER_PATTERN(IREE::VM::ListReserveOp, "iree_vm_list_reserve",
                         DenseSet<size_t>({0}), true);
   ADD_CONTAINER_PATTERN(IREE::VM::ListResizeOp, "iree_vm_list_resize",

diff --git a/compiler/src/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/buffer_ops.mlir b/compiler/src/iree/compiler/Dialect/VM/Conversion/VMToEmitC/test/buffer_ops.mlir
@@ -236,3 +236,16 @@ vm.module @my_module {
     vm.return
   }
 }
+
+// -----
+
+vm.module @my_module {
+  // CHECK-LABEL: @my_module_buffer_hash
+  vm.func @buffer_hash(%buf : !vm.buffer) {
+    // CHECK: %[[STATUS:.+]] = emitc.call_opaque "iree_vm_buffer_hash"
+    %c0 = vm.const.i64 0
+    %c10 = vm.const.i64 10
+    %v0 = vm.buffer.hash %buf, %c0, %c10 : !vm.buffer -> i64
+    vm.return
+  }
+}
diff --git a/compiler/src/iree/compiler/Dialect/VM/IR/VMOpcodesCore.td b/compiler/src/iree/compiler/Dialect/VM/IR/VMOpcodesCore.td
@@ -47,7 +47,7 @@ class VM_OPC_EnumAttr<string name, string enumName, string enumTag,
   string opcodeEnumTag = enumTag;
 }
 
-// Next available opcode: 0x84
+// Next available opcode: 0x85
 
 // Globals:
 def VM_OPC_GlobalLoadI32         : VM_OPC<0x00, "GlobalLoadI32">;
@@ -222,6 +222,7 @@ def VM_OPC_BufferFillI8          : VM_OPC<0x71, "BufferFillI8">;
 def VM_OPC_BufferFillI16         : VM_OPC<0x72, "BufferFillI16">;
 def VM_OPC_BufferFillI32         : VM_OPC<0x73, "BufferFillI32">;
 def VM_OPC_BufferFillI64         : VM_OPC<0x74, "BufferFillI64">;
+def VM_OPC_BufferHash            : VM_OPC<0x84, "BufferHash">;
 
 // Extension prefixes:
 def VM_OPC_PrefixExtF32          : VM_OPC<0xE0, "PrefixExtF32">;
@@ -380,6 +381,7 @@ def VM_CoreOpcodeAttr :
     VM_OPC_BufferFillI64,
     VM_OPC_BufferCopy,
     VM_OPC_BufferCompare,
+    VM_OPC_BufferHash,
 
     VM_OPC_Block,
 

diff --git a/compiler/src/iree/compiler/Dialect/VM/IR/VMOps.td b/compiler/src/iree/compiler/Dialect/VM/IR/VMOps.td
@@ -1598,6 +1598,38 @@ def VM_BufferStoreF64Op :
   let summary = [{64-bit floating-point store}];
 }
 
+def VM_BufferHashOp : VM_Op<"buffer.hash", [
+      DeclareOpInterfaceMethods<VM_SerializableOpInterface>,
+      MemoryEffects<[MemRead]>,
+    ]> {
+  let description = [{
+    Computes the SipHash-2-4 of the source buffer at the given offset for
+    |length| bytes using seed `0x0001020304...0e0f`.
+  }];
+
+  let arguments = (ins
+    VM_RefOf<VM_BufferType>:$source_buffer,
+    VM_BufferIndex:$source_offset,
+    VM_BufferIndex:$length
+  );
+  let results = (outs
+    I64:$result
+  );
+
+  let assemblyFormat = [{
+    $source_buffer `,` $source_offset `,` $length
+    attr-dict `:` type($source_buffer) `->` type($result)
+  }];
+
+  let encoding = [
+    VM_EncOpcode<VM_OPC_BufferHash>,
+    VM_EncOperand<"source_buffer", 0>,
+    VM_EncOperand<"source_offset", 1>,
+    VM_EncOperand<"length", 2>,
+    VM_EncResult<"result">,
+  ];
+}
+
 } // OpGroupBufferOps
 
 //===----------------------------------------------------------------------===//

diff --git a/runtime/src/iree/vm/buffer.c b/runtime/src/iree/vm/buffer.c
@@ -314,6 +314,91 @@ IREE_API_EXPORT iree_status_t iree_vm_buffer_write_elements(
   return iree_ok_status();
 }
 
+// Based on reference implementation from https://github.com/veorq/SipHash
+// By Jean-Philippe Aumasson and Daniel J. Bernstein. (CC0 Licensed)
+#define ROTL(x, b) (uint64_t)(((x) << (b)) | ((x) >> (64 - (b))))
+#define SIPROUND(v0, v1, v2, v3) \
+  v0 += v1;                      \
+  v1 = ROTL(v1, 13);             \
+  v1 ^= v0;                      \
+  v0 = ROTL(v0, 32);             \
+  v2 += v3;                      \
+  v3 = ROTL(v3, 16);             \
+  v3 ^= v2;                      \
+  v0 += v3;                      \
+  v3 = ROTL(v3, 21);             \
+  v3 ^= v0;                      \
+  v2 += v1;                      \
+  v1 = ROTL(v1, 17);             \
+  v1 ^= v2;                      \
+  v2 = ROTL(v2, 32)
+
+// Using SipHash-2-4.
+#ifndef cROUNDS
+#define cROUNDS 2
+#endif
+#ifndef dROUNDS
+#define dROUNDS 4
+#endif
+
+IREE_API_EXPORT iree_status_t iree_vm_buffer_hash(
+    const iree_vm_buffer_t* source_buffer, iree_host_size_t source_offset,
+    iree_host_size_t length, int64_t* out_result) {
+  IREE_TRACE_ZONE_BEGIN(z0);
+  IREE_ASSERT_ARGUMENT(source_buffer);
+
+  // Get the byte span for the source data.
+  iree_const_byte_span_t source_span = iree_const_byte_span_empty();
+  IREE_RETURN_AND_END_ZONE_IF_ERROR(
+      z0, iree_vm_buffer_map_ro(source_buffer, source_offset, length, 1,
+                                &source_span));
+  const uint8_t* source = source_span.data;
+  const uint8_t* end = source + source_span.data_length -
+                       (source_span.data_length % sizeof(uint64_t));
+  const int left = source_span.data_length & 7;
+  uint64_t hash = ((uint64_t)source_span.data_length) << 56;
+
+  // Using key = 0x000102030405060708090a0b0c0d0e0f
+  uint64_t v0 = UINT64_C(0x736f6d6570736575 ^ 0x0706050403020100);
+  uint64_t v1 = UINT64_C(0x646f72616e646f6d ^ 0x0f0e0d0c0b0a0908);
+  uint64_t v2 = UINT64_C(0x6c7967656e657261 ^ 0x0706050403020100);
+  uint64_t v3 = UINT64_C(0x7465646279746573 ^ 0x0f0e0d0c0b0a0908);
+  uint64_t m;
+
+  for (; source != end; source += 8) {
+    m = iree_unaligned_load_le_u64((const uint64_t*)source);
+    v3 ^= m;
+    for (int i = 0; i < cROUNDS; ++i) {
+      SIPROUND(v0, v1, v2, v3);
+    }
+    v0 ^= m;
+  }
+
+  uint64_t tmp = 0;
+  for (int l = left; l > 0; --l) {
+    tmp = tmp << 8;
+    tmp |= (uint64_t)source[l - 1];
+  }
+  hash |= tmp;
+  v3 ^= hash;
+
+  for (int i = 0; i < cROUNDS; ++i) {
+    SIPROUND(v0, v1, v2, v3);
+  }
+
+  v0 ^= hash;
+  v2 ^= 0xff;
+
+  for (int i = 0; i < dROUNDS; ++i) {
+    SIPROUND(v0, v1, v2, v3);
+  }
+
+  hash = v0 ^ v1 ^ v2 ^ v3;
+  *out_result = hash;
+  IREE_TRACE_ZONE_END(z0);
+  return iree_ok_status();
+}
+
 iree_status_t iree_vm_buffer_register_types(iree_vm_instance_t* instance) {
   static const iree_vm_ref_type_descriptor_t descriptor = {
       .destroy = iree_vm_buffer_destroy,

diff --git a/runtime/src/iree/vm/buffer.h b/runtime/src/iree/vm/buffer.h
@@ -199,6 +199,12 @@ IREE_API_EXPORT iree_status_t iree_vm_buffer_write_elements(
     iree_host_size_t target_offset, iree_host_size_t element_count,
     iree_host_size_t element_length);
 
+// Computes the SipHash-2-4 of the source buffer along the byte range starting
+// as |source_offset| of |length| bytes. Uses 128-bit key `0x00010203...0e0f`.
+IREE_API_EXPORT iree_status_t iree_vm_buffer_hash(
+    const iree_vm_buffer_t* source_buffer, iree_host_size_t source_offset,
+    iree_host_size_t length, int64_t* result);
+
 // Low-level helper for accessing a typed view of a buffer for read access.
 // The calling function must be safe to return from. Assumes buffer is non-null.
 // Prefer iree_vm_buffer_read_elements for larger reads.

diff --git a/runtime/src/iree/vm/bytecode/disassembler.c b/runtime/src/iree/vm/bytecode/disassembler.c
@@ -1024,6 +1024,27 @@ iree_status_t iree_vm_bytecode_disassemble_op(
       break;
     }
 
+    DISASM_OP(CORE, BufferHash) {
+      bool buffer_is_move;
+      uint16_t buffer_reg =
+          VM_ParseOperandRegRef("source_buffer", &buffer_is_move);
+      uint16_t offset_reg = VM_ParseOperandRegI64("source_offset");
+      uint16_t length_reg = VM_ParseOperandRegI64("length");
+      uint16_t result_reg = VM_ParseResultRegI64("result");
+      EMIT_I32_REG_NAME(result_reg);
+      IREE_RETURN_IF_ERROR(
+          iree_string_builder_append_cstring(b, " = vm.buffer.hash "));
+      EMIT_REF_REG_NAME(buffer_reg);
+      EMIT_OPTIONAL_VALUE_REF(&regs->ref[buffer_reg]);
+      IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(b, ", "));
+      EMIT_I64_REG_NAME(offset_reg);
+      EMIT_OPTIONAL_VALUE_I64(regs->i32[offset_reg]);
+      IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(b, ", "));
+      EMIT_I64_REG_NAME(length_reg);
+      EMIT_OPTIONAL_VALUE_I64(regs->i32[length_reg]);
+      break;
+    }
+
     //===------------------------------------------------------------------===//
     // Lists
     //===------------------------------------------------------------------===//

diff --git a/runtime/src/iree/vm/bytecode/dispatch.c b/runtime/src/iree/vm/bytecode/dispatch.c
@@ -1176,6 +1176,24 @@ static iree_status_t iree_vm_bytecode_dispatch(
       vm_buffer_store_i64_inline(buffer, offset, value);
     });
 
+    DISPATCH_OP(CORE, BufferHash, {
+      bool source_buffer_is_move;
+      iree_vm_ref_t* source_buffer_ref =
+          VM_DecOperandRegRef("source_buffer", &source_buffer_is_move);
+      iree_vm_buffer_t* source_buffer =
+          iree_vm_buffer_deref(*source_buffer_ref);
+      if (IREE_UNLIKELY(!source_buffer)) {
+        return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+                                "source_buffer is null");
+      }
+      iree_host_size_t source_offset =
+          VM_DecOperandRegI64HostSize("source_offset");
+      iree_host_size_t length = VM_DecOperandRegI64HostSize("length");
+      uint64_t* result = VM_DecResultRegI64("result");
+      IREE_RETURN_IF_ERROR(
+          iree_vm_buffer_hash(source_buffer, source_offset, length, result));
+    });
+
     //===------------------------------------------------------------------===//
     // Lists
     //===------------------------------------------------------------------===//