[HAL] Add hal.allocator.resolve_memory_properties (iree-org#21115)

Progress towards: iree-org#20856

Adds a new operation `hal.allocator.resolve_memory_properties` and a
pass that tries to resolve it using the topology attribute. The same
pass should aim to do the same thing for the `hal.allocator.select` op
in a follow up.

---------

Signed-off-by: default <[email protected]>
Signed-off-by: Thomas Ziereis <[email protected]>
Signed-off-by: keshavvinayak01 <[email protected]>

Author: ziereis

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

@@ -5,7 +5,6 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.h"
-
 #include "iree/compiler/Dialect/HAL/Analysis/Captures.h"
 #include "iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Utils.h"
 #include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
@@ -99,26 +98,23 @@ struct ResourceAllocOpPattern
     auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
     auto resourceType =
         cast<IREE::Stream::ResourceType>(allocOp.getResult().getType());
-    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
-    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-    if (failed(deriveAllowedResourceBufferBits(allocOp.getLoc(), resourceType,
-                                               memoryTypes, bufferUsage))) {
-      return failure();
-    }
 
-    Value memoryTypeOp =
-        rewriter.create<IREE::HAL::MemoryTypeOp>(allocOp.getLoc(), memoryTypes);
-    Value bufferUsageOp = rewriter.create<IREE::HAL::BufferUsageOp>(
-        allocOp.getLoc(), bufferUsage);
+    auto resolveOp =
+        rewriter.create<IREE::HAL::AllocatorResolveMemoryPropertiesOp>(
+            allocOp.getLoc(), rewriter.getI32Type(), rewriter.getI32Type(),
+            allocOp.getAffinity().value_or(nullptr),
+            static_cast<IREE::HAL::Lifetime>(resourceType.getLifetime()));
 
     // Lookup the appropriate allocator/queue for allocation based on the buffer
     // propreties.
     auto [allocator, queueAffinity] = lookupAllocatorAndQueueAffinityFor(
-        allocOp, memoryTypeOp, bufferUsageOp, rewriter);
+        allocOp, resolveOp.getMemoryTypes(), resolveOp.getBufferUsage(),
+        rewriter);
 
     rewriter.replaceOpWithNewOp<IREE::HAL::AllocatorAllocateOp>(
-        allocOp, bufferType, allocator, queueAffinity, memoryTypeOp,
-        bufferUsageOp, adaptor.getStorageSize());
+        allocOp, bufferType, allocator, queueAffinity,
+        resolveOp.getMemoryTypes(), resolveOp.getBufferUsage(),
+        adaptor.getStorageSize());
     return success();
   }
 };
@@ -134,22 +130,19 @@ struct ResourceAllocaOpPattern
     // Derive buffer propreties from the resource type.
     auto resourceType =
         cast<IREE::Stream::ResourceType>(allocaOp.getResult().getType());
-    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
-    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-    if (failed(deriveAllowedResourceBufferBits(loc, resourceType, memoryTypes,
-                                               bufferUsage))) {
-      return failure();
-    }
     auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
-    Value memoryTypeOp =
-        rewriter.create<IREE::HAL::MemoryTypeOp>(loc, memoryTypes);
-    Value bufferUsageOp =
-        rewriter.create<IREE::HAL::BufferUsageOp>(loc, bufferUsage);
 
     // Lookup the appropriate device/queue for allocation based on the buffer
     // propreties.
-    auto [device, queueAffinity] = lookupDeviceAndQueueAffinityFor(
-        allocaOp, memoryTypeOp, bufferUsageOp, rewriter);
+    auto resolveOp =
+        rewriter.create<IREE::HAL::AllocatorResolveMemoryPropertiesOp>(
+            loc, rewriter.getI32Type(), rewriter.getI32Type(),
+            allocaOp.getAffinity().value_or(nullptr),
+            static_cast<IREE::HAL::Lifetime>(resourceType.getLifetime()));
+
+    auto [device, queueAffinity] =
+        lookupDeviceAndQueueAffinityFor(allocaOp, resolveOp.getMemoryTypes(),
+                                        resolveOp.getBufferUsage(), rewriter);
 
     // Behavior flags.
     IREE::HAL::AllocaFlagBitfield flags = IREE::HAL::AllocaFlagBitfield::None;
@@ -167,7 +160,8 @@ struct ResourceAllocaOpPattern
     auto pool = rewriter.create<arith::ConstantIntOp>(loc, 0, 64);
     auto allocateOp = rewriter.create<IREE::HAL::DeviceQueueAllocaOp>(
         loc, bufferType, device, queueAffinity, waitFence, signalFence, pool,
-        memoryTypeOp, bufferUsageOp, adaptor.getStorageSize(), flags);
+        resolveOp.getMemoryTypes(), resolveOp.getBufferUsage(),
+        adaptor.getStorageSize(), flags);
 
     rewriter.replaceOp(allocaOp, {allocateOp.getResult(), signalFence});
     return success();
@@ -188,28 +182,24 @@ struct ResourceDeallocaOpPattern
     // prefer-origin mode.
     auto resourceType =
         cast<IREE::Stream::ResourceType>(deallocaOp.getOperand().getType());
-    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
-    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-    bool preferOrigin = deallocaOp.getPreferOrigin();
-    if (failed(deriveAllowedResourceBufferBits(loc, resourceType, memoryTypes,
-                                               bufferUsage))) {
-      preferOrigin = true;
-    }
 
-    Value memoryTypeOp =
-        rewriter.create<IREE::HAL::MemoryTypeOp>(loc, memoryTypes);
-    Value bufferUsageOp =
-        rewriter.create<IREE::HAL::BufferUsageOp>(loc, bufferUsage);
+    auto resolveOp =
+        rewriter.create<IREE::HAL::AllocatorResolveMemoryPropertiesOp>(
+            loc, rewriter.getI32Type(), rewriter.getI32Type(),
+            deallocaOp.getAffinity().value_or(nullptr),
+            static_cast<IREE::HAL::Lifetime>(resourceType.getLifetime()));
 
-    auto [device, queueAffinity] = lookupDeviceAndQueueAffinityFor(
-        deallocaOp, memoryTypeOp, bufferUsageOp, rewriter);
+    auto [device, queueAffinity] =
+        lookupDeviceAndQueueAffinityFor(deallocaOp, resolveOp.getMemoryTypes(),
+                                        resolveOp.getBufferUsage(), rewriter);
 
     // Gather wait/signal fence, which are optional.
     Value waitFence =
         getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
     Value signalFence = getOrCreateSignalFence(
         loc, device, deallocaOp.getResultTimepoint(), rewriter);
 
+    bool preferOrigin = deallocaOp.getPreferOrigin();
     // Route to the origin of the allocation (if available).
     IREE::HAL::DeallocaFlagBitfield flags =
         IREE::HAL::DeallocaFlagBitfield::None;
@@ -467,8 +457,9 @@ buildStorageAssertions(Location loc, Value buffer, StringAttr message,
                        OpBuilder &builder) {
   auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
   auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-  if (failed(deriveRequiredResourceBufferBits(loc, resourceType, memoryTypes,
-                                              bufferUsage))) {
+  if (failed(deriveRequiredResourceBufferBits(
+          loc, static_cast<IREE::HAL::Lifetime>(resourceType.getLifetime()),
+          memoryTypes, bufferUsage))) {
     return failure();
   }
 

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Utils.cpp

@@ -5,20 +5,9 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Utils.h"
-
 #include "iree/compiler/Dialect/HAL/Analysis/Captures.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
-#include "llvm/Support/CommandLine.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-
-static llvm::cl::opt<bool> clExternalResourcesMappable(
-    "iree-stream-external-resources-mappable",
-    llvm::cl::desc("Allocates external resources as host-visible and mappable. "
-                   "This can degrade performance and introduce allocation "
-                   "overhead and staging buffers for readback on the host "
-                   "should be managed by the calling application instead."),
-    llvm::cl::init(false));
 
 namespace mlir::iree_compiler {
 
@@ -199,36 +188,34 @@ IREE::HAL::CommandCategoryBitfield deriveCommandCategories(Region &region) {
 }
 
 LogicalResult
-deriveRequiredResourceBufferBits(Location loc,
-                                 IREE::Stream::ResourceType resourceType,
+deriveRequiredResourceBufferBits(Location loc, IREE::HAL::Lifetime lifetime,
                                  IREE::HAL::MemoryTypeBitfield &memoryTypes,
                                  IREE::HAL::BufferUsageBitfield &bufferUsage) {
   memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
   bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-  switch (resourceType.getLifetime()) {
+  switch (lifetime) {
   default:
-    return mlir::emitError(loc)
-           << "unsupported resource lifetime: "
-           << IREE::Stream::stringifyLifetime(resourceType.getLifetime());
-  case IREE::Stream::Lifetime::Constant:
+    return mlir::emitError(loc) << "unsupported resource lifetime: "
+                                << IREE::HAL::stringifyLifetime(lifetime);
+  case IREE::HAL::Lifetime::Constant:
     // Device local; copies required to get into external resources.
     memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
     bufferUsage =
         bufferUsage | IREE::HAL::BufferUsageBitfield::SharingImmutable;
     break;
-  case IREE::Stream::Lifetime::Variable:
+  case IREE::HAL::Lifetime::Variable:
     // Device local; copies required to get into external resources.
     memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
     break;
-  case IREE::Stream::Lifetime::External:
+  case IREE::HAL::Lifetime::External:
     // We only require device-visible for external buffers (as we don't today
     // do anything else with them on the host). They may be mappable for user
     // convenience. Ideally they would have been placed in device-local memory
     // but so long as they are device visible the program will execute
     // correctly.
     memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceVisible;
     break;
-  case IREE::Stream::Lifetime::Staging:
+  case IREE::HAL::Lifetime::Staging:
     // Host local; copies required to get into device resources.
     // We could vary this based on staging usage (upload/download) by
     // making it device-local|host-visible, but host-local means we have
@@ -238,7 +225,7 @@ deriveRequiredResourceBufferBits(Location loc,
     bufferUsage = bufferUsage | IREE::HAL::BufferUsageBitfield::Transfer |
                   IREE::HAL::BufferUsageBitfield::Mapping;
     break;
-  case IREE::Stream::Lifetime::Transient:
+  case IREE::HAL::Lifetime::Transient:
     // Device local; copies required to get into external resources.
     memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
     break;
@@ -251,41 +238,6 @@ deriveRequiredResourceBufferBits(Location loc,
   return success();
 }
 
-LogicalResult
-deriveAllowedResourceBufferBits(Location loc,
-                                IREE::Stream::ResourceType resourceType,
-                                IREE::HAL::MemoryTypeBitfield &memoryTypes,
-                                IREE::HAL::BufferUsageBitfield &bufferUsage) {
-  memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
-  bufferUsage = IREE::HAL::BufferUsageBitfield::None;
-  if (failed(deriveRequiredResourceBufferBits(loc, resourceType, memoryTypes,
-                                              bufferUsage))) {
-    return failure();
-  }
-  switch (resourceType.getLifetime()) {
-  default:
-    break;
-  case IREE::Stream::Lifetime::External:
-    if (clExternalResourcesMappable) {
-      // #yolo; these come from/go to outside the program.
-      // Today we assume they are device-local|host-visible just for
-      // practical purposes but that does not have to be true. We really
-      // want this to be something we analyze and handle on the edges
-      // (transferring devices/etc if needed).
-      memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal |
-                    IREE::HAL::MemoryTypeBitfield::HostVisible;
-      // NOTE: we may not map it but users may after they get them back.
-      // Another reason we should annotate this - having a buffer be
-      // mappable is potentially expensive (may get a 2nd copy in memory!).
-      bufferUsage = bufferUsage | IREE::HAL::BufferUsageBitfield::Mapping;
-    } else {
-      memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
-    }
-    break;
-  }
-  return success();
-}
-
 BindingTable::BindingTable(IREE::Stream::CmdExecuteOp executeOp,
                            ValueRange bufferValues, ValueRange bufferSizes,
                            IndexSet &indexSet) {

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Utils.h

@@ -7,10 +7,8 @@
 #ifndef IREE_COMPILER_DIALECT_HAL_CONVERSION_STREAMTOHAL_UTILS_H_
 #define IREE_COMPILER_DIALECT_HAL_CONVERSION_STREAMTOHAL_UTILS_H_
 
-#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
-#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
@@ -61,22 +59,10 @@ IREE::HAL::CommandCategoryBitfield deriveCommandCategories(Region &region);
 // The bits set here are those that must be set for the buffer to be used as the
 // buffer within the program with its defined resource lifetime.
 LogicalResult
-deriveRequiredResourceBufferBits(Location loc,
-                                 IREE::Stream::ResourceType resourceType,
+deriveRequiredResourceBufferBits(Location loc, IREE::HAL::Lifetime lifetime,
                                  IREE::HAL::MemoryTypeBitfield &memoryTypes,
                                  IREE::HAL::BufferUsageBitfield &bufferUsage);
 
-// Maps a resource type to the corresponding HAL memory types and buffer usage.
-// This will fail if the resource type is not directly mappable to HAL bits.
-// The bits set here represent the superset of required and allowed bits and
-// are useful for providing buffers back to users via the ABI that may need to
-// be used for more than just what the internal program requires.
-LogicalResult
-deriveAllowedResourceBufferBits(Location loc,
-                                IREE::Stream::ResourceType resourceType,
-                                IREE::HAL::MemoryTypeBitfield &memoryTypes,
-                                IREE::HAL::BufferUsageBitfield &bufferUsage);
-
 class BindingTable {
 public:
   BindingTable() = default;

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/resource_ops.mlir

@@ -4,10 +4,11 @@ util.global private @device : !hal.device
 
 // CHECK-LABEL: @resourceAlloc
 util.func public @resourceAlloc(%arg0: index) -> !stream.resource<transient> {
-  // CHECK: %[[MEMORY_TYPE:.+]] = hal.memory_type<"DeviceVisible|DeviceLocal"> : i32
-  // CHECK: %[[BUFFER_USAGE:.+]] = hal.buffer_usage<"{{.+}}Transfer{{.+}}Dispatch{{.+}}"> : i32
+  // CHECK: %[[MEMORY_TYPES:.+]], %[[BUFFER_USAGE:.+]] = hal.allocator.resolve_memory_properties
+  // CHECK-SAME: for(#hal.device.affinity<@device>)
+  // CHECK-SAME: lifetime(transient) : i32, i32
   // CHECK: %[[RET0:.+]] = hal.allocator.allocate
-  // CHECK-SAME: type(%[[MEMORY_TYPE]])
+  // CHECK-SAME: type(%[[MEMORY_TYPES]])
   // CHECK-SAME: usage(%[[BUFFER_USAGE]])
   // CHECK-SAME: : !hal.buffer{%arg0}
   %0 = stream.resource.alloc uninitialized on(#hal.device.affinity<@device>) : !stream.resource<transient>{%arg0}
@@ -22,16 +23,17 @@ util.global private @device : !hal.device
 // CHECK-LABEL: @resourceAlloca
 // CHECK-SAME: (%[[SIZE:.+]]: index)
 util.func public @resourceAlloca(%size: index) -> (!stream.resource<transient>, !stream.timepoint) {
-  // CHECK: %[[MEMORY_TYPE:.+]] = hal.memory_type<"DeviceVisible|DeviceLocal"> : i32
-  // CHECK: %[[BUFFER_USAGE:.+]] = hal.buffer_usage<"{{.+}}Transfer{{.+}}DispatchStorage"> : i32
+  // CHECK: %[[MEMORY_TYPES:.+]], %[[BUFFER_USAGE:.+]] = hal.allocator.resolve_memory_properties
+  // CHECK-SAME: for(#hal.device.affinity<@device>)
+  // CHECK-SAME: lifetime(transient) : i32, i32
   // CHECK: %[[WAIT_FENCE:.+]] = util.null : !hal.fence
   // CHECK: %[[SIGNAL_FENCE:.+]] = hal.fence.create
   // CHECK: %[[RET0:.+]] = hal.device.queue.alloca
   // CHECK-SAME: affinity(%c-1
   // CHECK-SAME: wait(%[[WAIT_FENCE]])
   // CHECK-SAME: signal(%[[SIGNAL_FENCE]])
   // CHECK-SAME: pool(%c0
-  // CHECK-SAME: type(%[[MEMORY_TYPE]])
+  // CHECK-SAME: type(%[[MEMORY_TYPES]])
   // CHECK-SAME: usage(%[[BUFFER_USAGE]])
   // CHECK-SAME: flags("None")
   // CHECK-SAME: : !hal.buffer{%[[SIZE]]}
@@ -47,15 +49,16 @@ util.global private @device : !hal.device
 // CHECK-LABEL: @resourceAllocaAwait
 // CHECK-SAME: (%[[SIZE:.+]]: index, %[[WAIT_FENCE:.+]]: !hal.fence)
 util.func public @resourceAllocaAwait(%size: index, %await_timepoint: !stream.timepoint) -> (!stream.resource<transient>, !stream.timepoint) {
-  // CHECK: %[[MEMORY_TYPE:.+]] = hal.memory_type<"DeviceVisible|DeviceLocal"> : i32
-  // CHECK: %[[BUFFER_USAGE:.+]] = hal.buffer_usage<"{{.+}}Transfer{{.+}}DispatchStorage"> : i32
+  // CHECK: %[[MEMORY_TYPES:.+]], %[[BUFFER_USAGE:.+]] = hal.allocator.resolve_memory_properties
+  // CHECK-SAME: for(#hal.device.affinity<@device>)
+  // CHECK-SAME: lifetime(transient) : i32, i32
   // CHECK: %[[SIGNAL_FENCE:.+]] = hal.fence.create
   // CHECK: %[[RET0:.+]] = hal.device.queue.alloca
   // CHECK-SAME: affinity(%c-1
   // CHECK-SAME: wait(%[[WAIT_FENCE]])
   // CHECK-SAME: signal(%[[SIGNAL_FENCE]])
   // CHECK-SAME: pool(%c0
-  // CHECK-SAME: type(%[[MEMORY_TYPE]])
+  // CHECK-SAME: type(%[[MEMORY_TYPES]])
   // CHECK-SAME: usage(%[[BUFFER_USAGE]])
   // CHECK-SAME: flags("None")
   // CHECK-SAME: : !hal.buffer{%[[SIZE]]}
@@ -72,16 +75,15 @@ util.global private @device_b : !hal.device
 // CHECK-LABEL: @resourceAllocaOptimal
 // CHECK-SAME: (%[[SIZE:.+]]: index)
 util.func public @resourceAllocaOptimal(%size: index) -> (!stream.resource<transient>, !stream.timepoint) {
-  // CHECK: %[[MEMORY_TYPE:.+]] = hal.memory_type<"DeviceVisible|DeviceLocal"> : i32
-  // CHECK: %[[BUFFER_USAGE:.+]] = hal.buffer_usage<"{{.+}}Transfer{{.+}}DispatchStorage"> : i32
+  // CHECK: %[[MEMORY_TYPES:.+]], %[[BUFFER_USAGE:.+]] = hal.allocator.resolve_memory_properties
   // CHECK: %[[DEVICE_A:.+]] = util.global.load immutable @device_a : !hal.device
   // CHECK: %[[AFFINITY_A:.+]] = arith.constant -1 : i64
   // CHECK: %[[DEVICE_B:.+]] = util.global.load immutable @device_b : !hal.device
   // CHECK: %[[AFFINITY_B:.+]] = arith.constant -1 : i64
   // CHECK: %[[DEVICE:.+]], %[[QUEUE_AFFINITY:.+]] = hal.allocator.select from([
   // CHECK: (%[[DEVICE_A]], %[[AFFINITY_A]] : !hal.device, i64),
   // CHECK: (%[[DEVICE_B]], %[[AFFINITY_B]] : !hal.device, i64)
-  // CHECK: ]) type(%[[MEMORY_TYPE]]) usage(%[[BUFFER_USAGE]]) : !hal.device, i64
+  // CHECK: ]) type(%[[MEMORY_TYPES]]) usage(%[[BUFFER_USAGE]]) : !hal.device, i64
   // CHECK: %[[WAIT_FENCE:.+]] = util.null : !hal.fence
   // CHECK: %[[SIGNAL_FENCE:.+]] = hal.fence.create device(%[[DEVICE]] : !hal.device)
   // CHECK: %[[POOL:.+]] = arith.constant 0 : i64
@@ -91,7 +93,7 @@ util.func public @resourceAllocaOptimal(%size: index) -> (!stream.resource<trans
   // CHECK-SAME: wait(%[[WAIT_FENCE]])
   // CHECK-SAME: signal(%[[SIGNAL_FENCE]])
   // CHECK-SAME: pool(%[[POOL]])
-  // CHECK-SAME: type(%[[MEMORY_TYPE]])
+  // CHECK-SAME: type(%[[MEMORY_TYPES]])
   // CHECK-SAME: usage(%[[BUFFER_USAGE]])
   // CHECK-SAME: flags("None")
   // CHECK-SAME: : !hal.buffer{%[[SIZE]]}