Merging the HAL rewrite branch to main

CMakeLists.txt

@@ -115,7 +115,6 @@ endif()
 # List of all HAL drivers to be built by default:
 set(IREE_ALL_HAL_DRIVERS
   DyLib
-  Metal
   VMLA
   Vulkan
 )
@@ -126,9 +125,6 @@ if(IREE_HAL_DRIVERS_TO_BUILD STREQUAL "all")
   # For Apple platforms we need to use Metal instead of Vulkan.
   if(APPLE)
     list(REMOVE_ITEM IREE_HAL_DRIVERS_TO_BUILD Vulkan)
-  else()
-    # And Metal isn't available on non-Apple platforms for sure.
-    list(REMOVE_ITEM IREE_HAL_DRIVERS_TO_BUILD Metal)
   endif()
 endif()
 message(STATUS "Building HAL drivers: ${IREE_HAL_DRIVERS_TO_BUILD}")

bindings/python/pyiree/rt/function_abi.cc

@@ -590,7 +590,8 @@ void FunctionAbi::RawUnpack(absl::Span<const Description> descs,
           throw RaiseValueError(
               "Could not deref result buffer view (wrong type?)");
         }
-        iree_hal_buffer* raw_buffer = iree_hal_buffer_view_buffer(buffer_view);
+        iree_hal_buffer_t* raw_buffer =
+            iree_hal_buffer_view_buffer(buffer_view);
         if (!raw_buffer) {
           throw RaiseValueError("Could not deref result buffer (wrong type?)");
         }
@@ -675,10 +676,10 @@ void FunctionAbi::AllocateResults(absl::Span<const Description> descs,
             kScalarTypeToHalElementType[static_cast<unsigned>(
                 desc.scalar.type)]);
         iree_hal_buffer_view_t* buffer_view;
-        CheckApiStatus(iree_hal_buffer_view_create(
-                           raw_buffer, dims.data(), dims.size(), element_type,
-                           iree_allocator_system(), &buffer_view),
-                       "Error allocating buffer_view");
+        CheckApiStatus(
+            iree_hal_buffer_view_create(raw_buffer, dims.data(), dims.size(),
+                                        element_type, &buffer_view),
+            "Error allocating buffer_view");
         iree_hal_buffer_release(raw_buffer);
         iree_vm_ref_t buffer_view_ref =
             iree_hal_buffer_view_move_ref(buffer_view);
@@ -763,10 +764,10 @@ void FunctionAbi::PackBuffer(const RawSignatureParser::Description& desc,
   absl::InlinedVector<int, 5> dims(py_view.ndim);
   std::copy(py_view.shape, py_view.shape + py_view.ndim, dims.begin());
   iree_hal_buffer_view_t* buffer_view;
-  CheckApiStatus(iree_hal_buffer_view_create(
-                     raw_buffer, dims.data(), dims.size(), element_type,
-                     iree_allocator_system(), &buffer_view),
-                 "Error allocating buffer_view");
+  CheckApiStatus(
+      iree_hal_buffer_view_create(raw_buffer, dims.data(), dims.size(),
+                                  element_type, &buffer_view),
+      "Error allocating buffer_view");
   iree_hal_buffer_release(raw_buffer);
   iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view);
   CheckApiStatus(iree_vm_list_push_ref_move(f_args.raw_ptr(), &buffer_view_ref),

bindings/python/pyiree/rt/hal.cc

@@ -24,15 +24,15 @@ namespace {
 
 class HalMappedMemory {
  public:
-  HalMappedMemory(iree_hal_mapped_memory_t mapped_memory,
+  HalMappedMemory(iree_hal_buffer_mapping_t mapped_memory,
                   iree_hal_buffer_view_t* bv)
       : mapped_memory_(mapped_memory), bv_(bv) {
     iree_hal_buffer_view_retain(bv_);
   }
   ~HalMappedMemory() {
     if (bv_) {
       iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bv_);
-      IREE_CHECK_OK(iree_hal_buffer_unmap(buffer, &mapped_memory_));
+      iree_hal_buffer_unmap_range(&mapped_memory_);
       iree_hal_buffer_view_release(bv_);
     }
   }
@@ -44,10 +44,10 @@ class HalMappedMemory {
   static HalMappedMemory Create(HalBufferView& bv) {
     iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bv.raw_ptr());
     iree_device_size_t byte_length = iree_hal_buffer_byte_length(buffer);
-    iree_hal_mapped_memory_t mapped_memory;
-    CheckApiStatus(iree_hal_buffer_map(buffer, IREE_HAL_MEMORY_ACCESS_READ,
-                                       0 /* element_offset */, byte_length,
-                                       &mapped_memory),
+    iree_hal_buffer_mapping_t mapped_memory;
+    CheckApiStatus(iree_hal_buffer_map_range(
+                       buffer, IREE_HAL_MEMORY_ACCESS_READ,
+                       0 /* element_offset */, byte_length, &mapped_memory),
                    "Could not map memory");
     return HalMappedMemory(mapped_memory, bv.raw_ptr());
   }
@@ -81,7 +81,7 @@ class HalMappedMemory {
   }
 
  private:
-  iree_hal_mapped_memory_t mapped_memory_;
+  iree_hal_buffer_mapping_t mapped_memory_;
   iree_hal_buffer_view_t* bv_;
 };
 
@@ -167,9 +167,6 @@ void SetupHalBindings(pybind11::module m) {
   py::class_<HalMappedMemory>(m, "MappedMemory", py::buffer_protocol())
       .def_buffer(&HalMappedMemory::ToBufferInfo);
   py::class_<HalBuffer>(m, "HalBuffer")
-      .def_static("allocate_heap", &HalBuffer::AllocateHeapBuffer,
-                  py::arg("memory_type"), py::arg("usage"),
-                  py::arg("allocation_size"))
       .def("fill_zero", &HalBuffer::FillZero, py::arg("byte_offset"),
            py::arg("byte_length"))
       .def("create_view", &HalBuffer::CreateView, py::arg("shape"),

bindings/python/pyiree/rt/hal.h

@@ -92,18 +92,6 @@ class HalBufferView
 
 class HalBuffer : public ApiRefCounted<HalBuffer, iree_hal_buffer_t> {
  public:
-  static HalBuffer AllocateHeapBuffer(int32_t memory_type, int32_t usage,
-                                      iree_host_size_t allocation_size) {
-    iree_hal_buffer_t* buffer = nullptr;
-    CheckApiStatus(
-        iree_hal_heap_buffer_allocate(
-            static_cast<iree_hal_memory_type_t>(memory_type),
-            static_cast<iree_hal_buffer_usage_t>(usage), allocation_size,
-            iree_allocator_system(), iree_allocator_system(), &buffer),
-        "Error allocating heap buffer");
-    return HalBuffer::CreateRetained(buffer);
-  }
-
   iree_device_size_t byte_length() const {
     return iree_hal_buffer_byte_length(raw_ptr());
   }
@@ -121,7 +109,7 @@ class HalBuffer : public ApiRefCounted<HalBuffer, iree_hal_buffer_t> {
         IREE_HAL_ELEMENT_TYPE_NONE, element_size * 8);
     CheckApiStatus(
         iree_hal_buffer_view_create(raw_ptr(), shape.s.data(), shape.s.size(),
-                                    element_type, iree_allocator_system(), &bv),
+                                    element_type, &bv),
         "Error creating buffer view");
     return HalBufferView::CreateRetained(bv);
   }

bindings/python/pyiree/rt/hal_test.py

@@ -25,27 +25,6 @@ def testEnums(self):
     logging.info("MemoryType: %s", rt.MemoryType)
     logging.info("HOST_VISIBLE: %s", int(rt.MemoryType.HOST_VISIBLE))
 
-  def testAllocateHeap(self):
-    b = rt.HalBuffer.allocate_heap(memory_type=int(rt.MemoryType.HOST_LOCAL),
-                                   usage=int(rt.BufferUsage.ALL),
-                                   allocation_size=4096)
-    self.assertIsNot(b, None)
-    b.fill_zero(0, 4096)
-    shape = rt.Shape([1, 1024])
-    unused_bv = b.create_view(shape, 4)
-
-  def testStrideCalculation(self):
-    b = rt.HalBuffer.allocate_heap(memory_type=int(rt.MemoryType.HOST_LOCAL),
-                                   usage=int(rt.BufferUsage.ALL),
-                                   allocation_size=4096)
-    self.assertIsNot(b, None)
-    b.fill_zero(0, 4096)
-    shape = rt.Shape([16, 1, 8, 4, 2])
-    bv = b.create_view(shape, 4)
-    self.assertEqual(
-        np.array(bv.map()).strides,
-        (1 * 8 * 4 * 2 * 4, 8 * 4 * 2 * 4, 4 * 2 * 4, 2 * 4, 4))
-
 
 if __name__ == "__main__":
   absltest.main()

bindings/python/pyiree/rt/host_types.cc

@@ -119,15 +119,14 @@ class PyMappedMemory {
     }
   };
 
-  PyMappedMemory(Description desc, iree_hal_mapped_memory_t mapped_memory,
+  PyMappedMemory(Description desc, iree_hal_buffer_mapping_t mapped_memory,
                  HalBuffer buffer)
       : desc_(std::move(desc)),
         mapped_memory_(mapped_memory),
         buf_(std::move(buffer)) {}
   ~PyMappedMemory() {
     if (buf_) {
-      CheckApiStatus(iree_hal_buffer_unmap(buf_.raw_ptr(), &mapped_memory_),
-                     "Error unmapping memory");
+      iree_hal_buffer_unmap_range(&mapped_memory_);
     }
   }
   PyMappedMemory(PyMappedMemory&& other)
@@ -139,8 +138,8 @@ class PyMappedMemory {
                                               HalBuffer buffer) {
     iree_device_size_t byte_length =
         iree_hal_buffer_byte_length(buffer.raw_ptr());
-    iree_hal_mapped_memory_t mapped_memory;
-    CheckApiStatus(iree_hal_buffer_map(
+    iree_hal_buffer_mapping_t mapped_memory;
+    CheckApiStatus(iree_hal_buffer_map_range(
                        buffer.raw_ptr(), IREE_HAL_MEMORY_ACCESS_READ,
                        0 /* element_offset */, byte_length, &mapped_memory),
                    "Could not map memory");
@@ -160,7 +159,7 @@ class PyMappedMemory {
 
  private:
   Description desc_;
-  iree_hal_mapped_memory_t mapped_memory_;
+  iree_hal_buffer_mapping_t mapped_memory_;
   HalBuffer buf_;
 };
 

build_tools/bazel/iree.bazelrc

@@ -159,26 +159,54 @@ build:macos_clang --per_file_copt=tensorflow,iree_tf_compiler@-Wno-range-loop-an
 # https://github.com/google/sanitizers/wiki/AddressSanitizer
 ###############################################################################
 
-# Turn on asan. Some toolchains make use of the asan feature and we'll directly
-# set the appropriate opts.
+# ASAN (address sanitizer)
+# https://clang.llvm.org/docs/AddressSanitizer.html
+build:asan --config=sanitizer
 build:asan --features=asan
 build:asan --copt=-fsanitize=address
+build:asan --copt=-fsanitize-address-use-after-scope
 build:asan --linkopt=-fsanitize=address
+build:asan --cc_output_directory_tag=asan
+build:asan --copt=-DADDRESS_SANITIZER
+
+# MSAN (memory sanitizer)
+# https://clang.llvm.org/docs/MemorySanitizer.html
+build:msan --config=sanitizer
+build:msan --features=msan
+build:msan --copt=-fsanitize=memory
+build:msan --copt=-fsanitize-memory-track-origins
+build:msan --linkopt=-fsanitize=memory
+build:msan --cc_output_directory_tag=msan
+build:msan --copt=-DMEMORY_SANITIZER
+
+# TSAN (thread sanitizer)
+# https://clang.llvm.org/docs/ThreadSanitizer.html
+build:tsan --config=sanitizer
+build:tsan --features=tsan
+build:tsan --copt=-fsanitize=thread
+build:tsan --linkopt=-fsanitize=thread
+build:tsan --cc_output_directory_tag=tsan
+build:tsan --copt=-DTHREAD_SANITIZER
+
+# UBSAN (undefined behavior sanitizer)
+# https://clang.llvm.org/docs/UndefinedBehaviorSanitizer.html
+build:ubsan --config=sanitizer
+build:ubsan --features=ubsan
+build:ubsan --copt=-fsanitize=undefined
+build:ubsan --linkopt=-fsanitize=undefined
+build:ubsan --linkopt=-lubsan
+build:ubsan --cc_output_directory_tag=ubsan
 
 # Don't strip debug info
-build:asan --strip=never
+build:sanitizer --strip=never
 # Ignore settings of `linkopts = ["-static"]` which can screw up the sanitizer.
 # We don't use this in IREE (that's what linkstatic is for), but it could show
 # up in dependencies.
-build:asan --force_ignore_dash_static
-# asan tests tend to take longer, so increase the timeouts
-build:asan --test_timeout=120,600,1800,-1
-# Make the outputs easy to find
-build:asan --cc_output_directory_tag=asan
+build:sanitizer --force_ignore_dash_static
+# sanitizer tests tend to take longer, so increase the timeouts
+build:sanitizer --test_timeout=120,600,1800,-1
 # Get better stack traces
-build:asan --copt=-fno-omit-frame-pointer
-# This macro define is used by absl
-build:asan --copt=-DADDRESS_SANITIZER
+build:sanitizer --copt=-fno-omit-frame-pointer
 
 ###############################################################################
 # Architecture specific options

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -61,7 +61,7 @@
     "@renderdoc_api//:renderdoc_app": ["renderdoc_api::renderdoc_app"],
     "@pffft": ["pffft"],
     "@cpuinfo//:cpuinfo": ["cpuinfo"],
-    "@half//:half": ["half"],
+    "@half//:includes": ["half::includes"],
     "@vulkan_memory_allocator//:impl_header_only": ["vulkan_memory_allocator"],
 }
 

build_tools/cmake/iree_macros.cmake

@@ -306,10 +306,10 @@ function(iree_add_test_environment_properties TEST_NAME)
   #
   # Tests which only depend on a compiler target backend or a runtime HAL
   # driver, but not both, should generally use a different method of filtering.
-  if(NOT ${IREE_TARGET_BACKEND_VULKAN-SPIRV} OR NOT ${IREE_HAL_DRIVER_VULKAN})
+  if(NOT "${IREE_TARGET_BACKEND_VULKAN-SPIRV}" OR NOT "${IREE_HAL_DRIVER_VULKAN}")
     set_property(TEST ${TEST_NAME} APPEND PROPERTY ENVIRONMENT "IREE_VULKAN_DISABLE=1")
   endif()
-  if(NOT ${IREE_TARGET_BACKEND_DYLIB-LLVM-AOT} OR NOT ${IREE_HAL_DRIVER_DYLIB})
+  if(NOT "${IREE_TARGET_BACKEND_DYLIB-LLVM-AOT}" OR NOT "${IREE_HAL_DRIVER_DYLIB}")
     set_property(TEST ${TEST_NAME} APPEND PROPERTY ENVIRONMENT "IREE_LLVMAOT_DISABLE=1")
   endif()
 endfunction()