Adding flags to most HAL methods and extending existing ones to i64. (#20368)

**This is a breaking HAL change and will require recompilation of
VMFBs.** There _should_ be a useful error indicating that emitted if
attempting to load an existing VMFB.

There's been TODOs in the HAL imports for awhile waiting for a breakage
and #20240 does that. To prevent more PRs like that this adds flag
arguments to most HAL methods so that we can add new flags in the
future. Some existing flags were i32 and have been changed to i64 for
consistency.

Of note:
* `hal.device.queue.execute` split to `hal.device.queue.barrier` instead
of relying on empty command buffer list
* `hal.command_buffer.advise_buffer` added, but not yet in IR
* `hal.command_buffer.fill_buffer` max pattern length changed to i64 to
match `hal.device.queue.fill`

Author: benvanik

compiler/plugins/input/Torch/InputConversion/FuncConversion.cpp

@@ -541,8 +541,9 @@ void createCoarseFencesSyncWrapper(StringRef syncFunctionName,
           .getResults();
 
   // Wait forever for signal.
-  rewriter.create<IREE::HAL::FenceAwaitOp>(loc, rewriter.getI32Type(),
-                                           timeoutMillis, signalFence);
+  rewriter.create<IREE::HAL::FenceAwaitOp>(
+      loc, rewriter.getI32Type(), timeoutMillis,
+      IREE::HAL::WaitFlagBitfield::None, signalFence);
 
   rewriter.create<IREE::Util::ReturnOp>(loc, callResults);
 }

compiler/src/iree/compiler/Bindings/Native/Transforms/WrapEntryPoints.cpp

@@ -259,9 +259,9 @@ createImportWrapperFunc(IREE::ABI::InvocationModel invocationModel,
   if (hasSideEffects && signalFence) {
     auto timeoutMillis =
         entryBuilder.create<arith::ConstantIntOp>(importOp.getLoc(), -1, 32);
-    entryBuilder.create<IREE::HAL::FenceAwaitOp>(importOp.getLoc(),
-                                                 entryBuilder.getI32Type(),
-                                                 timeoutMillis, signalFence);
+    entryBuilder.create<IREE::HAL::FenceAwaitOp>(
+        importOp.getLoc(), entryBuilder.getI32Type(), timeoutMillis,
+        IREE::HAL::WaitFlagBitfield::None, signalFence);
   }
 
   // Marshal results.

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

@@ -4,6 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "iree/compiler/Dialect/HAL/Conversion/HALToVM/Patterns.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
 #include "iree/compiler/Dialect/VM/Conversion/ImportUtils.h"
@@ -137,13 +138,16 @@ class CommandBufferFillBufferOpConversion
     auto patternLengthConst = rewriter.createOrFold<mlir::arith::ConstantIntOp>(
         op.getLoc(), patternLengthBytes, 32);
     Value pattern = op.getPattern();
-    if (patternBitWidth < 32) {
+    if (patternBitWidth < 64) {
       pattern = rewriter.createOrFold<arith::ExtUIOp>(
-          op.getLoc(), rewriter.getIntegerType(32), pattern);
+          op.getLoc(), rewriter.getIntegerType(64), pattern);
     }
     callOperands.push_back(pattern);
     callOperands.push_back(patternLengthConst);
 
+    callOperands.push_back(
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter));
+
     auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>(
         op, SymbolRefAttr::get(importOp), importType.getResults(),
         callOperands);
@@ -183,7 +187,9 @@ class CommandBufferUpdateBufferOpConversion
         castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(),
                          rewriter),
         castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter),
-        targetBufferSlot};
+        targetBufferSlot,
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
+    };
     auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>(
         op, SymbolRefAttr::get(importOp), importType.getResults(),
         callOperands);
@@ -226,6 +232,7 @@ class CommandBufferCopyBufferOpConversion
         castToImportType(adaptor.getTargetOffset(), rewriter.getI64Type(),
                          rewriter),
         castToImportType(adaptor.getLength(), rewriter.getI64Type(), rewriter),
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
     };
     auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>(
         op, SymbolRefAttr::get(importOp), importType.getResults(),
@@ -341,21 +348,14 @@ class CommandBufferDispatchOpConversion
     auto i64Type = rewriter.getI64Type();
     Value zeroI32 = rewriter.create<IREE::VM::ConstI32ZeroOp>(op.getLoc());
 
-    auto flags = adaptor.getFlagsAttr()
-                     ? rewriter
-                           .create<IREE::VM::ConstI64Op>(
-                               op.getLoc(), adaptor.getFlagsAttr().getInt())
-                           .getResult()
-                     : rewriter.create<IREE::VM::ConstI64ZeroOp>(op.getLoc())
-                           .getResult();
     SmallVector<Value, 8> callOperands = {
         adaptor.getCommandBuffer(),
         adaptor.getExecutable(),
         castToImportType(adaptor.getEntryPoint(), i32Type, rewriter),
         castToImportType(adaptor.getWorkgroupX(), i32Type, rewriter),
         castToImportType(adaptor.getWorkgroupY(), i32Type, rewriter),
         castToImportType(adaptor.getWorkgroupZ(), i32Type, rewriter),
-        flags,
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
     };
     SmallVector<int16_t, 5> segmentSizes = {
         /*command_buffer=*/-1,
@@ -421,21 +421,14 @@ class CommandBufferDispatchIndirectOpConversion
 
     auto [workgroupsBufferSlot, workgroupsBuffer] =
         splitBufferSlot(op.getLoc(), adaptor.getWorkgroupsBuffer(), rewriter);
-    auto flags = adaptor.getFlagsAttr()
-                     ? rewriter
-                           .create<IREE::VM::ConstI64Op>(
-                               op.getLoc(), adaptor.getFlagsAttr().getInt())
-                           .getResult()
-                     : rewriter.create<IREE::VM::ConstI64ZeroOp>(op.getLoc())
-                           .getResult();
     SmallVector<Value, 8> callOperands = {
         adaptor.getCommandBuffer(),
         adaptor.getExecutable(),
         castToImportType(adaptor.getEntryPoint(), i32Type, rewriter),
         workgroupsBufferSlot,
         workgroupsBuffer,
         castToImportType(adaptor.getWorkgroupsOffset(), i64Type, rewriter),
-        flags,
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
     };
     SmallVector<int16_t, 5> segmentSizes = {
         /*command_buffer=*/-1,

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

@@ -4,6 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "iree/compiler/Dialect/HAL/Conversion/HALToVM/Patterns.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/VM/Conversion/ImportUtils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -134,8 +135,6 @@ class DeviceQueueFillOpConversion
     auto patternLength = rewriter.create<IREE::VM::ConstI32Op>(
         op.getLoc(),
         llvm::divideCeil(op.getPattern().getType().getIntOrFloatBitWidth(), 8));
-    auto flags =
-        rewriter.create<IREE::VM::ConstI64Op>(op.getLoc(), op.getFlags());
     std::array<Value, 10> callOperands = {
         adaptor.getDevice(),
         castToImportType(adaptor.getQueueAffinity(), i64Type, rewriter),
@@ -146,7 +145,7 @@ class DeviceQueueFillOpConversion
         castToImportType(adaptor.getLength(), i64Type, rewriter),
         castToImportType(adaptor.getPattern(), i64Type, rewriter),
         patternLength,
-        flags,
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
     };
     auto callOp = rewriter.replaceOpWithNewOp<IREE::VM::CallOp>(
         op, SymbolRefAttr::get(importOp), importType.getResults(),
@@ -177,19 +176,23 @@ class DeviceQueueExecuteIndirectOpConversion
     auto importType = importOp.getFunctionType();
     auto i64Type = rewriter.getI64Type();
 
+    Value queueAffinity =
+        castToImportType(adaptor.getQueueAffinity(), i64Type, rewriter);
     SmallVector<Value, 8> callOperands = {
         adaptor.getDevice(),
-        castToImportType(adaptor.getQueueAffinity(), i64Type, rewriter),
+        queueAffinity,
         adaptor.getWaitFence(),
         adaptor.getSignalFence(),
         adaptor.getCommandBuffer(),
+        getFlagsI64(op.getLoc(), adaptor.getFlagsAttr(), rewriter),
     };
     SmallVector<int16_t, 5> segmentSizes = {
         /*device=*/-1,
         /*queue_affinity=*/-1,
         /*wait_fence=*/-1,
         /*signal_fence=*/-1,
         /*command_buffer=*/-1,
+        /*flags=*/-1,
         /*bindings=*/
         static_cast<int16_t>(adaptor.getBindingBuffers().size()),
     };
@@ -239,6 +242,8 @@ void populateHALDeviceToVMPatterns(MLIRContext *context,
       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::DeviceQueueBarrierOp>>(
+      context, importSymbols, typeConverter, "hal.device.queue.barrier");
   patterns.insert<VMImportOpConversion<IREE::HAL::DeviceQueueExecuteOp>>(
       context, importSymbols, typeConverter, "hal.device.queue.execute");
   patterns.insert<DeviceQueueExecuteIndirectOpConversion>(

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

@@ -121,9 +121,8 @@ class ExecutableCreateOpConversion
         createOp.getLoc(), adaptor.getConstants(), rewriter);
 
     SmallVector<Value, 8> callOperands = {
-        adaptor.getDevice(),
-        executableFormatStr,
-        rodataOp,
+        adaptor.getDevice(), adaptor.getQueueAffinity(),
+        executableFormatStr, rodataOp,
         constantBuffer,
     };
     auto importType = importOp.getFunctionType();

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

@@ -10,6 +10,13 @@
 
 namespace mlir::iree_compiler {
 
+Value getFlagsI64(Location loc, IntegerAttr flagsAttr, OpBuilder &builder) {
+  return flagsAttr
+             ? builder.create<IREE::VM::ConstI64Op>(loc, flagsAttr.getInt())
+                   .getResult()
+             : builder.create<IREE::VM::ConstI64ZeroOp>(loc).getResult();
+}
+
 extern void populateHALAllocatorToVMPatterns(MLIRContext *context,
                                              SymbolTable &importSymbols,
                                              TypeConverter &typeConverter,

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/Patterns.h

@@ -23,6 +23,9 @@ void populateHALToVMPatterns(MLIRContext *context, SymbolTable &importSymbols,
 Value createPackedConstantBuffer(Location loc, ValueRange constantValues,
                                  OpBuilder &builder);
 
+// Returns an i64 value initialized with the bits of |flagsAttr| or 0.
+Value getFlagsI64(Location loc, IntegerAttr flagsAttr, OpBuilder &builder);
+
 } // namespace mlir::iree_compiler
 
 #endif // IREE_COMPILER_DIALECT_HAL_CONVERSION_HALTOVM_PATTERNS_H_

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

@@ -3,11 +3,11 @@
 // CHECK-LABEL: @channel_create
 //  CHECK-SAME: (%[[DEVICE:.+]]: !vm.ref<!hal.device>, %[[AFFINITY:.+]]: i64, %[[ID:.+]]: !vm.buffer, %[[GROUP:.+]]: !vm.buffer, %[[RANK:.+]]: i32, %[[COUNT:.+]]: i32) -> !vm.ref<!hal.channel>
 util.func public @channel_create(%device: !hal.device, %affinity: i64, %id: !util.buffer, %group: !util.buffer, %rank: i32, %count: i32) -> !hal.channel {
-  // CHECK: %[[FLAGS:.+]] = vm.const.i32.zero
+  // CHECK: %[[FLAGS:.+]] = vm.const.i64.zero
   // CHECK: %[[CHANNEL:.+]] = vm.call @hal.channel.create(%[[DEVICE]], %[[AFFINITY]], %[[FLAGS]], %[[ID]], %[[GROUP]], %[[RANK]], %[[COUNT]])
   %channel = hal.channel.create device(%device : !hal.device)
                               affinity(%affinity)
-                                 flags(0)
+                                 flags("None")
                                     id(%id)
                                  group(%group)
                                   rank(%rank)
@@ -21,12 +21,12 @@ util.func public @channel_create(%device: !hal.device, %affinity: i64, %id: !uti
 // CHECK-LABEL: @channel_split
 //  CHECK-SAME: (%[[BASE_CHANNEL:.+]]: !vm.ref<!hal.channel>, %[[COLOR:.+]]: i32, %[[KEY:.+]]: i32)
 util.func public @channel_split(%base_channel: !hal.channel, %color: i32, %key: i32) -> !hal.channel {
-  // CHECK: %[[FLAGS:.+]] = vm.const.i32.zero
+  // CHECK: %[[FLAGS:.+]] = vm.const.i64.zero
   // CHECK: %[[SPLIT_CHANNEL:.+]] = vm.call @hal.channel.split(%[[BASE_CHANNEL]], %[[COLOR]], %[[KEY]], %[[FLAGS]])
   %split_channel = hal.channel.split<%base_channel : !hal.channel>
                                color(%color)
                                  key(%key)
-                               flags(0) : !hal.channel
+                               flags("None") : !hal.channel
   // CHECK: vm.return %[[SPLIT_CHANNEL]]
   util.return %split_channel : !hal.channel
 }