Reworking constant upload with a HAL file API. (iree-org#14665)

This changes the compiler-generated map/staging IR into file reads that
are managed by the runtime HAL targets. In this initial version only one
type of file is supported that wraps host memory and exposes it via the
file API. For targets that don't natively support file streaming a
utility for asynchronously streaming files using an iree_loop_t is
available and all targets currently use that. In the future targets like
CUDA/Metal/Direct3D12 can use cuFile/MTLIOCommandBuffer/DirectStorage to
directly stream file contents from/to disk.

Future changes will expand the iree_hal_file_t vtable with what can be
reliably implemented or guarded with feature flags. When any form of
asynchronous loop is implemented (such as a task-system fiber loop)
we'll be able to perform overlapped copies to reduce total transfer
latency by letting the CPU and GPU do their respective staging
operations concurrently. Today only synchronous loops are used so it's
fully serialized.

Now that the compiler doesn't emit staging with
iree_allocator_allocate_buffer's initial_data that can be removed in
iree-org#14605 to remove the primary use of the existing buffer_transfer util.
The remaining places using it can be switched to memory file streaming
instead such as numpy IO and support much larger content.

This is a breaking HAL change and all existing VMFBs will fail to load
on newer runtimes.

CPU and CUDA targets benefit from this as should all other backends
though some may need to have zero-copy paths added - they'll at least
benefit from streaming and bounded staging buffer sizes. Vulkan has been
optimized a bit and the follow-on sparse bindings support improves it
further.

Vulkan dedicated GPU w/ 1GB model before this change (1024MB + 1024MB of
staging, even when mmapped):

![image](https://github.com/openxla/iree/assets/75337/62b4b144-efe6-4116-8192-3f7bbfbbfaeb)

Vulkan dedicated GPU w/ 1GB model now (1024MB + ~64MB of staging used as
streamed from a memory mapped file):

![image](https://github.com/openxla/iree/assets/75337/4e44a186-9c11-4d83-951b-b6bc1ad6e3f2)

Vulkan integrated GPU w/ 1GB model now (1024MB of host memory imported
and used directly with no copies):

![image](https://github.com/openxla/iree/assets/75337/0d036ede-0aac-45be-b069-6bb8d6e06cb6)

Progress on iree-org#14607.

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertAllocatorOps.cpp

@@ -13,21 +13,20 @@ namespace mlir {
 namespace iree_compiler {
 namespace {
 
-class AllocatorAllocateInitializedOpConversion
-    : public OpConversionPattern<IREE::HAL::AllocatorAllocateInitializedOp> {
+class AllocatorAllocateOpConversion
+    : public OpConversionPattern<IREE::HAL::AllocatorAllocateOp> {
 public:
-  AllocatorAllocateInitializedOpConversion(TypeConverter &typeConverter,
-                                           MLIRContext *context,
-                                           SymbolTable &importSymbols)
+  AllocatorAllocateOpConversion(TypeConverter &typeConverter,
+                                MLIRContext *context,
+                                SymbolTable &importSymbols)
       : OpConversionPattern(typeConverter, context) {
-    importOp = importSymbols.lookup<IREE::VM::ImportOp>(
-        "hal.allocator.allocate.initialized");
+    importOp =
+        importSymbols.lookup<IREE::VM::ImportOp>("hal.allocator.allocate");
     assert(importOp);
   }
 
   LogicalResult
-  matchAndRewrite(IREE::HAL::AllocatorAllocateInitializedOp op,
-                  OpAdaptor adaptor,
+  matchAndRewrite(IREE::HAL::AllocatorAllocateOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>(
         op, importOp.getName(),
@@ -36,14 +35,13 @@ class AllocatorAllocateInitializedOpConversion
         },
         ArrayRef<Value>{
             adaptor.getAllocator(),
+            castToImportType(adaptor.getQueueAffinity(), rewriter.getI64Type(),
+                             rewriter),
             rewriter.createOrFold<IREE::VM::ConstI32Op>(
                 op.getLoc(), op.getMemoryTypesAttr().getInt()),
             rewriter.createOrFold<IREE::VM::ConstI32Op>(
                 op.getLoc(), op.getBufferUsageAttr().getInt()),
-            adaptor.getSource(),
-            castToImportType(adaptor.getOffset(), rewriter.getI64Type(),
-                             rewriter),
-            castToImportType(adaptor.getLength(), rewriter.getI64Type(),
+            castToImportType(adaptor.getResultSize(), rewriter.getI64Type(),
                              rewriter),
         });
     copyImportAttrs(importOp, callOp);
@@ -54,19 +52,18 @@ class AllocatorAllocateInitializedOpConversion
   mutable IREE::VM::ImportOp importOp;
 };
 
-class AllocatorTryMapOpConversion
-    : public OpConversionPattern<IREE::HAL::AllocatorTryMapOp> {
+class AllocatorImportOpConversion
+    : public OpConversionPattern<IREE::HAL::AllocatorImportOp> {
 public:
-  AllocatorTryMapOpConversion(TypeConverter &typeConverter,
+  AllocatorImportOpConversion(TypeConverter &typeConverter,
                               MLIRContext *context, SymbolTable &importSymbols)
       : OpConversionPattern(typeConverter, context) {
-    importOp = importSymbols.lookup<IREE::VM::ImportOp>(
-        "hal.allocator.map.byte_buffer");
+    importOp = importSymbols.lookup<IREE::VM::ImportOp>("hal.allocator.import");
     assert(importOp);
   }
 
   LogicalResult
-  matchAndRewrite(IREE::HAL::AllocatorTryMapOp op, OpAdaptor adaptor,
+  matchAndRewrite(IREE::HAL::AllocatorImportOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto callOp = rewriter.create<IREE::VM::CallOp>(
         op.getLoc(), importOp.getName(),
@@ -76,6 +73,8 @@ class AllocatorTryMapOpConversion
         ArrayRef<Value>{
             adaptor.getAllocator(),
             rewriter.createOrFold<IREE::VM::ConstI32Op>(op.getLoc(), /*try=*/1),
+            castToImportType(adaptor.getQueueAffinity(), rewriter.getI64Type(),
+                             rewriter),
             rewriter.createOrFold<IREE::VM::ConstI32Op>(
                 op.getLoc(), op.getMemoryTypesAttr().getInt()),
             rewriter.createOrFold<IREE::VM::ConstI32Op>(
@@ -88,9 +87,9 @@ class AllocatorTryMapOpConversion
         });
     copyImportAttrs(importOp, callOp);
     auto result = callOp.getResults().front();
-    auto didMap = rewriter.create<IREE::VM::CmpNZRefOp>(
+    auto didImport = rewriter.create<IREE::VM::CmpNZRefOp>(
         op.getLoc(), rewriter.getI32Type(), result);
-    rewriter.replaceOp(op, {didMap, result});
+    rewriter.replaceOp(op, {didImport, result});
     return success();
   }
 
@@ -104,11 +103,9 @@ void populateHALAllocatorToVMPatterns(MLIRContext *context,
                                       SymbolTable &importSymbols,
                                       TypeConverter &typeConverter,
                                       RewritePatternSet &patterns) {
-  patterns.insert<VMImportOpConversion<IREE::HAL::AllocatorAllocateOp>>(
-      context, importSymbols, typeConverter, "hal.allocator.allocate");
-  patterns.insert<AllocatorAllocateInitializedOpConversion>(
-      typeConverter, context, importSymbols);
-  patterns.insert<AllocatorTryMapOpConversion>(typeConverter, context,
+  patterns.insert<AllocatorAllocateOpConversion>(typeConverter, context,
+                                                 importSymbols);
+  patterns.insert<AllocatorImportOpConversion>(typeConverter, context,
                                                importSymbols);
 }
 

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertDeviceOps.cpp

@@ -129,6 +129,10 @@ void populateHALDeviceToVMPatterns(MLIRContext *context,
       context, importSymbols, typeConverter, "hal.device.queue.alloca");
   patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueDeallocaOp>>(
       context, importSymbols, typeConverter, "hal.device.queue.dealloca");
+  patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueReadOp>>(
+      context, importSymbols, typeConverter, "hal.device.queue.read");
+  patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueWriteOp>>(
+      context, importSymbols, typeConverter, "hal.device.queue.write");
   patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueExecuteOp>>(
       context, importSymbols, typeConverter, "hal.device.queue.execute");
   patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueFlushOp>>(

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/ConvertExperimentalOps.cpp

@@ -17,6 +17,8 @@ void populateHALExperimentalToVMPatterns(MLIRContext *context,
                                          RewritePatternSet &patterns) {
   patterns.insert<VMImportOpConversion<IREE::HAL::ExSharedDeviceOp>>(
       context, importSymbols, typeConverter, "hal.ex.shared_device");
+  patterns.insert<VMImportOpConversion<IREE::HAL::ExFileFromMemoryOp>>(
+      context, importSymbols, typeConverter, "hal.ex.file.from_memory");
 }
 
 } // namespace iree_compiler

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/allocator_ops.mlir

@@ -2,22 +2,28 @@
 
 // CHECK-LABEL: vm.func private @allocatorAllocate
 func.func @allocatorAllocate(%arg0 : !hal.allocator) -> !hal.buffer {
-  // CHECK: %[[SIZE:.+]] = vm.const.i64 1024
-  %c1024 = arith.constant 1024 : index
-  // CHECK: %ref = vm.call @hal.allocator.allocate(%arg0, %c70, %c3075, %[[SIZE]]) : (!vm.ref<!hal.allocator>, i32, i32, i64) -> !vm.ref<!hal.buffer>
-  %0 = hal.allocator.allocate<%arg0 : !hal.allocator> type("HostLocal") usage("DispatchStorage|Transfer") : !hal.buffer{%c1024}
+  // CHECK-DAG: %[[SIZE:.+]] = vm.const.i64 1024
+  %size = arith.constant 1024 : index
+  // CHECK-DAG: %[[AFFINITY:.+]] = vm.const.i64 -1
+  %affinity = arith.constant -1 : i64
+  // CHECK: %ref = vm.call @hal.allocator.allocate(%arg0, %[[AFFINITY]], %c70, %c3075, %[[SIZE]]) : (!vm.ref<!hal.allocator>, i64, i32, i32, i64) -> !vm.ref<!hal.buffer>
+  %0 = hal.allocator.allocate<%arg0 : !hal.allocator> affinity(%affinity) type("HostLocal") usage("DispatchStorage|Transfer") : !hal.buffer{%size}
   return %0 : !hal.buffer
 }
 
 // -----
 
-// CHECK-LABEL: vm.func private @allocatorMapByteBuffer
-func.func @allocatorMapByteBuffer(%arg0 : !hal.allocator, %arg1 : !util.buffer) -> !hal.buffer {
+// CHECK-LABEL: vm.func private @allocatorImport
+func.func @allocatorImport(%arg0 : !hal.allocator, %arg1 : !util.buffer) -> (i1, !hal.buffer) {
   // CHECK-DAG: %[[OFFSET:.+]] = vm.const.i64 128
   %offset = arith.constant 128 : index
   // CHECK-DAG: %[[LENGTH:.+]] = vm.const.i64 256
   %length = arith.constant 256 : index
-  // CHECK: = vm.call @hal.allocator.allocate.initialized(%arg0, %c6, %c3, %arg1, %[[OFFSET]], %[[LENGTH]]) : (!vm.ref<!hal.allocator>, i32, i32, !vm.buffer, i64, i64) -> !vm.ref<!hal.buffer>
-  %buffer = hal.allocator.allocate.initialized<%arg0 : !hal.allocator> source(%arg1 : !util.buffer)[%offset, %length] type("HostVisible|HostCoherent") usage("Transfer") : !hal.buffer
-  return %buffer : !hal.buffer
+  // CHECK-DAG: %[[AFFINITY:.+]] = vm.const.i64 -1
+  %affinity = arith.constant -1 : i64
+  // CHECK: %[[IMPORTED:.+]] = vm.call @hal.allocator.import(%arg0, %c1, %[[AFFINITY]], %c6, %c3, %arg1, %[[OFFSET]], %[[LENGTH]]) : (!vm.ref<!hal.allocator>, i32, i64, i32, i32, !vm.buffer, i64, i64) -> !vm.ref<!hal.buffer>
+  %did_import, %buffer = hal.allocator.import<%arg0 : !hal.allocator> source(%arg1 : !util.buffer)[%offset, %length] affinity(%affinity) type("HostVisible|HostCoherent") usage("Transfer") : i1, !hal.buffer
+  // CHECK: %[[DID_IMPORT:.+]] = vm.cmp.nz.ref %[[IMPORTED]]
+  // CHECK: return %[[DID_IMPORT]], %[[IMPORTED]]
+  return %did_import, %buffer : i1, !hal.buffer
 }

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/test/device_ops.mlir

@@ -143,6 +143,41 @@ func.func @device_queue_dealloca(
 
 // -----
 
+// CHECK-LABEL: @device_queue_read
+func.func @device_queue_read(
+    // CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>, %[[AFFINITY:.+]]: i64,
+    %device: !hal.device, %affinity: i64,
+    // CHECK-SAME:  %[[WAIT_FENCE:.+]]: !vm.ref<!hal.fence>, %[[SIGNAL_FENCE:.+]]: !vm.ref<!hal.fence>,
+    %wait_fence: !hal.fence, %signal_fence: !hal.fence,
+    // CHECK-SAME:  %[[SOURCE_FILE:.+]]: !vm.ref<!hal.file>,
+    %source_file: !hal.file,
+    // CHECK-SAME:  %[[TARGET_BUFFER:.+]]: !vm.ref<!hal.buffer>)
+    %target_buffer: !hal.buffer) {
+  // CHECK-DAG: %[[SOURCE_OFFSET:.+]] = vm.const.i64 100
+  %source_offset = arith.constant 100 : i64
+  // CHECK-DAG: %[[TARGET_OFFSET:.+]] = vm.const.i64 200
+  %target_offset = arith.constant 200 : index
+  // CHECK-DAG: %[[LENGTH:.+]] = vm.const.i64 300
+  %length = arith.constant 300 : index
+  // CHECK-DAG: %[[FLAGS:.+]] = vm.const.i32.zero
+  // CHECK: vm.call @hal.device.queue.read(
+  // CHECK-SAME: %[[DEVICE]], %[[AFFINITY]],
+  // CHECK-SAME: %[[WAIT_FENCE]], %[[SIGNAL_FENCE]],
+  // CHECK-SAME: %[[SOURCE_FILE]], %[[SOURCE_OFFSET]],
+  // CHECK-SAME: %[[TARGET_BUFFER]], %[[TARGET_OFFSET]],
+  // CHECK-SAME: %[[LENGTH]], %[[FLAGS]])
+  hal.device.queue.read<%device : !hal.device>
+      affinity(%affinity)
+      wait(%wait_fence) signal(%signal_fence)
+      source(%source_file : !hal.file)[%source_offset]
+      target(%target_buffer : !hal.buffer)[%target_offset]
+      length(%length)
+      flags(0)
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @device_queue_execute
 func.func @device_queue_execute(
     // CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>, %[[AFFINITY:.+]]: i64,

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/BUILD.bazel

@@ -33,6 +33,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Transforms",
     ],
 )

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/CMakeLists.txt

@@ -23,6 +23,7 @@ iree_cc_library(
     MLIRFuncDialect
     MLIRIR
     MLIRPass
+    MLIRSCFDialect
     MLIRTransforms
     iree::compiler::Dialect::HAL::Conversion
     iree::compiler::Dialect::HAL::IR