Merge branch 'main' into google-to-main

SUBMODULE_VERSIONS.txt

@@ -14,7 +14,7 @@ d8c7ee00a687ac369e62e2032514a93a9b413502 third_party/pybind11
 685f86471e9d26b3eb7676695a2e2cefb4551ae9 third_party/spirv_cross
 f8bf11a0253a32375c32cad92c841237b96696c0 third_party/spirv_headers
 75e42f8f26b75fe5e8af8461152fd1389e1c8229 third_party/tensorflow
-8732f0e94e4e41049a43029202bda94d7b4e85da third_party/tracy
+e35d02186d00eff26ec6c698331a0767a28a51b0 third_party/tracy
 9bd3f561bcee3f01d22912de10bb07ce4e23d378 third_party/vulkan_headers
 3528e2aed3e8808f33e1e7d63eeb1560456a605a third_party/vulkan_memory_allocator
 54e82ec439552018d0f26c60b0fd6cc9e36b42a8 third_party/wasm-micro-runtime

build_tools/bazel/build_core.sh

@@ -48,6 +48,9 @@ declare -a test_env_args=(
 declare -a default_build_tag_filters=("-nokokoro")
 declare -a default_test_tag_filters=("-nokokoro" "-driver=metal")
 
+# The VK_KHR_shader_float16_int8 extension is optional prior to Vulkan 1.2.
+# We test on SwiftShader, which does not support this extension.
+default_test_tag_filters+=("-vulkan_uses_vk_khr_shader_float16_int8")
 # CUDA CI testing disabled until we setup a target for it.
 default_test_tag_filters+=("-driver=cuda")
 

build_tools/buildkite/cmake/android/arm64-v8a/pipeline.yml

@@ -59,7 +59,9 @@ steps:
       - "tar xzf build-artifacts.tgz"
       - "find build-android/ -name '*.cmake' -exec sed -i \"s!\\$IREE_DOCKER_WORKDIR/!\\$PWD/!g\" {} \\;"
       - "cd build-android/"
-      - "ctest --output-on-failure"
+      # vulkan tests using khr_shader_float16_int8 are failing on pixel4.
+      # Disabling it until we identify the root cause.
+      - "ctest --output-on-failure --label-exclude \"^vulkan_uses_vk_khr_shader_float16_int8\\$\""
     agents:
       - "android-soc=snapdragon-855"
       - "queue=test-android"

build_tools/cmake/test.sh

@@ -30,6 +30,9 @@ export IREE_VULKAN_DISABLE=${IREE_VULKAN_DISABLE:-0}
 export IREE_LLVMAOT_DISABLE=${IREE_LLVMAOT_DISABLE:-0}
 # CUDA is off by default.
 export IREE_CUDA_DISABLE=${IREE_CUDA_DISABLE:-1}
+# The VK_KHR_shader_float16_int8 extension is optional prior to Vulkan 1.2.
+# We test on SwiftShader, which does not support this extension.
+export IREE_VULKAN_F16_DISABLE=${IREE_VULKAN_F16_DISABLE:-1}
 
 # Tests to exclude by label. In addition to any custom labels (which are carried
 # over from Bazel tags), every test should be labeled with the directory it is
@@ -62,6 +65,9 @@ fi
 if [[ "${IREE_CUDA_DISABLE?}" == 1 ]]; then
   label_exclude_args+=("^driver=cuda$")
 fi
+if [[ "${IREE_VULKAN_F16_DISABLE?}" == 1 ]]; then
+  label_exclude_args+=("^vulkan_uses_vk_khr_shader_float16_int8$")
+fi
 
 # Join on "|"
 label_exclude_regex="($(IFS="|" ; echo "${label_exclude_args[*]?}"))"

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh

@@ -71,4 +71,4 @@ export CTEST_PARALLEL_LEVEL=${CTEST_PARALLEL_LEVEL:-$(nproc)}
 echo "Testing with CTest"
 ctest --output-on-failure \
    --tests-regex "^integrations/tensorflow/|^bindings/python/" \
-   --label-exclude "^nokokoro$"
+   --label-exclude "^nokokoro$|^vulkan_uses_vk_khr_shader_float16_int8$"

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build.sh

@@ -44,6 +44,7 @@ echo "Initializing submodules"
 echo "Building with cmake"
 ./build_tools/cmake/clean_build.sh
 
+export IREE_VULKAN_F16_DISABLE=0
 export IREE_CUDA_DISABLE=0
 echo "Testing with ctest"
 ./build_tools/cmake/test.sh

iree/base/tracing.h

@@ -450,12 +450,18 @@ enum {
 
 #define IREE_TRACE_ALLOC(ptr, size) ___tracy_emit_memory_alloc(ptr, size, 0)
 #define IREE_TRACE_FREE(ptr) ___tracy_emit_memory_free(ptr, 0)
+#define IREE_TRACE_ALLOC_NAMED(name, ptr, size) \
+  ___tracy_emit_memory_alloc_named(ptr, size, 0, name)
+#define IREE_TRACE_FREE_NAMED(name, ptr) \
+  ___tracy_emit_memory_free_named(ptr, 0, name)
 
 #endif  // IREE_TRACING_FEATURE_ALLOCATION_CALLSTACKS
 
 #else
 #define IREE_TRACE_ALLOC(ptr, size)
 #define IREE_TRACE_FREE(ptr)
+#define IREE_TRACE_ALLOC_NAMED(name, ptr, size)
+#define IREE_TRACE_FREE_NAMED(name, ptr)
 #endif  // IREE_TRACING_FEATURE_ALLOCATION_TRACKING
 
 #if defined(__cplusplus) && \

iree/compiler/Dialect/Flow/IR/FlowOps.cpp

@@ -19,6 +19,7 @@
 #include "iree/compiler/Dialect/Shape/IR/Builders.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/CommandLine.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
@@ -33,6 +34,12 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/RegionUtils.h"
 
+static llvm::cl::opt<int> clInlineConstantByteLength(
+    "iree-flow-inline-constants-max-byte-length",
+    llvm::cl::desc("Maximum byte-length of constant that can be inlined into a "
+                   "dispatch region"),
+    llvm::cl::init(256));
+
 namespace mlir {
 namespace iree_compiler {
 namespace IREE {
@@ -831,7 +838,8 @@ static bool canDispatchRegionContainOp(Operation *op) {
       uint64_t estimatedByteLength =
           (shapedType.getNumElements() * shapedType.getElementTypeBitWidth()) /
           8;
-      return denseAttr.isSplat() || estimatedByteLength <= 256;  // or whatever
+      return denseAttr.isSplat() ||
+             estimatedByteLength <= clInlineConstantByteLength;
     } else if (constantType.isIntOrIndexOrFloat()) {
       return true;
     }

iree/hal/command_buffer.c

@@ -26,14 +26,15 @@ IREE_HAL_API_RETAIN_RELEASE(command_buffer);
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_command_buffer_create(
     iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(device);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
   *out_command_buffer = NULL;
   IREE_TRACE_ZONE_BEGIN(z0);
   iree_status_t status =
       IREE_HAL_VTABLE_DISPATCH(device, iree_hal_device, create_command_buffer)(
-          device, mode, command_categories, out_command_buffer);
+          device, mode, command_categories, queue_affinity, out_command_buffer);
   IREE_TRACE_ZONE_END(z0);
   return status;
 }

iree/hal/command_buffer.h

@@ -64,6 +64,22 @@ enum iree_hal_command_category_e {
 };
 typedef uint32_t iree_hal_command_category_t;
 
+// A bitmask indicating affinity for a submission to use a particular set of
+// queues.
+//
+// Upon submission the queue is selected based on the flags set in
+// |command_categories| and the |queue_affinity|. As the number of available
+// queues can vary the |queue_affinity| is used to hash into the available
+// queues for the required categories. For example if 2 queues support transfer
+// commands and the affinity is 5 the resulting queue could be index hash(5)=1.
+// The affinity can thus be treated as just a way to indicate whether two
+// submissions must be placed on to the same queue. Note that the exact hashing
+// function is implementation dependent.
+typedef uint64_t iree_hal_queue_affinity_t;
+
+// Specifies that any queue may be selected.
+#define IREE_HAL_QUEUE_AFFINITY_ANY ((iree_hal_queue_affinity_t)(-1))
+
 // Bitfield specifying which execution stage a barrier should start/end at.
 //
 // Maps to VkPipelineStageFlagBits.
@@ -193,9 +209,14 @@ typedef struct iree_hal_command_buffer_s iree_hal_command_buffer_t;
 
 // Creates a command buffer ready to begin recording, possibly reusing an
 // existing one from the |device| pool.
+//
+// |queue_affinity| specifies the device queues the command buffer may be
+// submitted to. The queue affinity provided to iree_hal_device_queue_submit
+// must match or be a subset of the |queue_affinity|.
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_command_buffer_create(
     iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Retains the given |command_buffer| for the caller.

iree/hal/cts/command_buffer_test.cc

@@ -41,7 +41,8 @@ TEST_P(CommandBufferTest, Create) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   EXPECT_TRUE((iree_hal_command_buffer_allowed_categories(command_buffer) &
                IREE_HAL_COMMAND_CATEGORY_DISPATCH) ==
@@ -54,7 +55,8 @@ TEST_P(CommandBufferTest, BeginEnd) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -66,7 +68,8 @@ TEST_P(CommandBufferTest, SubmitEmpty) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -81,7 +84,8 @@ TEST_P(CommandBufferTest, FillBufferWithRepeatedBytes) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   iree_hal_buffer_t* device_buffer;
   IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer(
@@ -141,7 +145,8 @@ TEST_P(CommandBufferTest, CopyWholeBuffer) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   // Create and fill a host buffer.
   iree_hal_buffer_t* host_buffer;
@@ -192,7 +197,8 @@ TEST_P(CommandBufferTest, CopySubBuffer) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   iree_hal_buffer_t* device_buffer;
   IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer(

iree/hal/cts/event_test.cc

@@ -36,7 +36,8 @@ TEST_P(EventTest, SignalAndReset) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_signal_event(
@@ -60,10 +61,12 @@ TEST_P(EventTest, SubmitWithChainedCommandBuffers) {
   iree_hal_command_buffer_t* command_buffer_2;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer_1));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer_1));
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer_2));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer_2));
 
   // First command buffer signals the event when it completes.
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer_1));

iree/hal/cts/semaphore_submission_test.cc

@@ -57,7 +57,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitAndSignal) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -94,7 +95,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitWithWait) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
 
@@ -142,7 +144,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitWithMultipleSemaphores) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));

iree/hal/cuda/cuda_device.c

@@ -155,10 +155,12 @@ static iree_hal_allocator_t* iree_hal_cuda_device_allocator(
 static iree_status_t iree_hal_cuda_device_create_command_buffer(
     iree_hal_device_t* base_device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
   return iree_hal_cuda_graph_command_buffer_allocate(
-      &device->context_wrapper, mode, command_categories, out_command_buffer);
+      &device->context_wrapper, mode, command_categories, queue_affinity,
+      out_command_buffer);
 }
 
 static iree_status_t iree_hal_cuda_device_create_descriptor_set(
@@ -218,8 +220,9 @@ static iree_status_t iree_hal_cuda_device_create_semaphore(
 
 static iree_status_t iree_hal_cuda_device_queue_submit(
     iree_hal_device_t* base_device,
-    iree_hal_command_category_t command_categories, uint64_t queue_affinity,
-    iree_host_size_t batch_count, const iree_hal_submission_batch_t* batches) {
+    iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
+    const iree_hal_submission_batch_t* batches) {
   iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
   for (int i = 0; i < batch_count; i++) {
     for (int j = 0; j < batches[i].command_buffer_count; j++) {

iree/hal/cuda/graph_command_buffer.c

@@ -28,6 +28,7 @@ typedef struct {
   iree_hal_cuda_context_wrapper_t* context;
   iree_hal_command_buffer_mode_t mode;
   iree_hal_command_category_t allowed_categories;
+  iree_hal_queue_affinity_t queue_affinity;
   CUgraph graph;
   CUgraphExec exec;
   // Keep track of the last node added to the command buffer as we are currently
@@ -52,6 +53,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     iree_hal_cuda_context_wrapper_t* context,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(context);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
@@ -72,6 +74,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     command_buffer->context = context;
     command_buffer->mode = mode;
     command_buffer->allowed_categories = command_categories;
+    command_buffer->queue_affinity = queue_affinity;
     command_buffer->graph = graph;
     command_buffer->exec = NULL;
     command_buffer->last_node = NULL;

iree/hal/cuda/graph_command_buffer.h

@@ -29,6 +29,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     iree_hal_cuda_context_wrapper_t* context,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Returns the native cuda graph associated to the command buffer.

iree/hal/device.c

@@ -42,7 +42,7 @@ iree_hal_device_allocator(iree_hal_device_t* device) {
 
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_device_queue_submit(
     iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-    uint64_t queue_affinity, iree_host_size_t batch_count,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
     const iree_hal_submission_batch_t* batches) {
   IREE_ASSERT_ARGUMENT(device);
   IREE_ASSERT_ARGUMENT(!batch_count || batches);

iree/hal/device.h

@@ -171,7 +171,7 @@ iree_hal_device_allocator(iree_hal_device_t* device);
 // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/vkQueueSubmit.html
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_device_queue_submit(
     iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-    uint64_t queue_affinity, iree_host_size_t batch_count,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
     const iree_hal_submission_batch_t* batches);
 
 // Blocks the caller until the semaphores reach or exceed the specified payload
@@ -242,6 +242,7 @@ typedef struct {
   iree_status_t(IREE_API_PTR* create_command_buffer)(
       iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
       iree_hal_command_category_t command_categories,
+      iree_hal_queue_affinity_t queue_affinity,
       iree_hal_command_buffer_t** out_command_buffer);
 
   iree_status_t(IREE_API_PTR* create_descriptor_set)(
@@ -276,7 +277,7 @@ typedef struct {
 
   iree_status_t(IREE_API_PTR* queue_submit)(
       iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-      uint64_t queue_affinity, iree_host_size_t batch_count,
+      iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
       const iree_hal_submission_batch_t* batches);
 
   iree_status_t(IREE_API_PTR* wait_semaphores_with_deadline)(