Bump IREE to iree-org/iree@c37c680ae6e71 (#1350)

Contains xrt/xrt-lite changes due to IREE HAL runtime API changes.
Also linalg.Conv* named ops are now lowered as linalg.generic - so a few
changes to deal with the same has been added. One of the Conv tests has
been disabled for now.

Signed-off-by: Abhishek Varma <[email protected]>

Author: Abhishek-Varma

build_tools/ci/cpu_comparison/run.py

@@ -2472,22 +2472,24 @@ def __init__(self):
                 )
             )
 
+        # TODO: Named Conv* ops are lowered as linalg.generic ops causing issues in the expected
+        #       lowering config. Re-enable after fixing those.
         # Depthwise convolution tests:
-        depthwise_map = {
-            "conv_type": "depthwise_conv_2d_nhwc_hwc",
-            "N": 1,
-            "IH": 14,
-            "IC": 64,
-            "KH": 3,
-            "input_element_type": "i32",
-            "output_element_type": "i32",
-        }
-        generator = ConvolutionMlirGenerator(**depthwise_map)
-        self.register(
-            ConvolutionFromTemplate(
-                generator, TestParams(tile_pipeline="conv-decompose", n_repeats=2)
-            )
-        )
+        # depthwise_map = {
+        #     "conv_type": "depthwise_conv_2d_nhwc_hwc",
+        #     "N": 1,
+        #     "IH": 14,
+        #     "IC": 64,
+        #     "KH": 3,
+        #     "input_element_type": "i32",
+        #     "output_element_type": "i32",
+        # }
+        # generator = ConvolutionMlirGenerator(**depthwise_map)
+        # self.register(
+        #     ConvolutionFromTemplate(
+        #         generator, TestParams(tile_pipeline="conv-decompose", n_repeats=2)
+        #     )
+        # )
 
         # Softmax tests:
         # Note: The error tolerance for npu4 is higher than that for npu1_4col.

compiler/plugins/target/AMD-AIE/iree-amd-aie/Test/samples/conv2d_nhwc_objectfifo_e2e.mlir

@@ -8,7 +8,7 @@ func.func @conv_2d_nhwc_hwcf(%arg0: tensor<2x14x14x32xi32>, %arg1: tensor<3x3x32
   return %2 : tensor<2x12x12x64xi32>
 }
 
-// CHECK-LABEL: hal.executable.export public @conv_2d_nhwc_hwcf_dispatch_0_conv_2d_nhwc_hwcf_2x12x12x64x3x3x32_i32
+// CHECK-LABEL: hal.executable.export public @conv_2d_nhwc_hwcf_dispatch_0_conv_2x12x12x64x3x3x32_i32
 //       CHECK:    aie.device(npu1_4col)
 //       CHECK:    aie.shim_dma_allocation
 //       CHECK:    aie.shim_dma_allocation
@@ -24,7 +24,7 @@ func.func @conv_2d_nhwc_hwcf_q(%arg0: tensor<2x14x14x32xi8>, %arg1: tensor<3x3x3
   return %2 : tensor<2x12x12x64xi32>
 }
 
-// CHECK-LABEL: hal.executable.export public @conv_2d_nhwc_hwcf_q_dispatch_0_conv_2d_nhwc_hwcf_2x12x12x64x3x3x32_i8xi8xi32
+// CHECK-LABEL: hal.executable.export public @conv_2d_nhwc_hwcf_q_dispatch_0_conv_2x12x12x64x3x3x32_i8xi8xi32
 //       CHECK:    aie.device(npu1_4col)
 //       CHECK:    aie.shim_dma_allocation
 //       CHECK:    aie.shim_dma_allocation

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEPackAndTranspose.cpp

@@ -6,6 +6,7 @@
 
 #include "iree-amd-aie/IR/AMDAIEAttrs.h"
 #include "iree-amd-aie/Transforms/Passes.h"
+#include "iree-amd-aie/Transforms/Utils/AMDAIEUtils.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
@@ -62,7 +63,7 @@ void AMDAIEPackAndTransposePass::runOnOperation() {
   linalg::LinalgOp linalgOp;
   funcOp->walk([&](linalg::LinalgOp op) {
     if (linalg::isaContractionOpInterface(op) ||
-        isa<linalg::ConvolutionOpInterface>(op.getOperation())) {
+        isConvOp(dyn_cast<linalg::GenericOp>(op.getOperation()))) {
       linalgOp = op;
       return WalkResult::interrupt();
     }

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/KernelDispatch.cpp

@@ -703,7 +703,7 @@ static LogicalResult setRootConfigForConvDecomposePipeline(
   SmallVector<int64_t> packingSizes;
 
   // [N, OH, OW, OC, KH, KW, IC].
-  if (isa<linalg::Conv2DNhwcHwcfOp>(linalgOp)) {
+  if (isConv2DNhwcHwcfOp(linalgOp.getOperation())) {
     // The goal is to pack the input image and kernel as follows, when moving
     // from L2 to L1 (example where there are 32 input channels):
     // Image: memref<1x3x6x32xbf16> ->  memref<1x3x4x6x8xbf16>
@@ -720,25 +720,32 @@ static LogicalResult setRootConfigForConvDecomposePipeline(
     tileSizeLevel2 = {0, 0, 0, 0, 1, 1, 1, 0, 0};
   }
 
-  // [N, OC, OH, OW, IC, KH, KW]
-  else if (isa<linalg::Conv2DNchwFchwOp>(linalgOp)) {
-    // The matmul reduction dimension is the input channel (IC) dimension.
-    // For Conv2DNhwcHwcfOp, this dimension is already the inner-most dimension
-    // of the input image, and the penultimate dimension of the kernel --
-    // exactly what we want. For Conv2DNchwFchwOp, can the tensor dimensions be
-    // permuted in DMA to get them in the correct positions? For the image
-    // tensor, only if H*W is a nice power of 2 (DMA constraint). For kernels,
-    // it requires h*w is a nice power of 2 -- unlikely, we typically have
-    // h=w=3. The dimension permutations will therefore often therefore need to
-    // be done on the core. We leave this for future work, the expectation for
-    // now is that models have been transformed at a high level to avoid
-    // channel-first convolutions.
-    return linalgOp.emitError(
-        "Only channel-last convolution supported currently.");
-  }
-
-  // [N, OH, OW, C, KW, HW]
-  else if (isa<linalg::DepthwiseConv2DNhwcHwcOp>(linalgOp)) {
+  // // [N, OC, OH, OW, IC, KH, KW]
+  // TODO(avarma): Currently since we anyway don't support Conv2DNchwFchwOp, the
+  // following check has been disabled. else if
+  // (isa<linalg::Conv2DNchwFchwOp>(linalgOp)) {
+  //   // The matmul reduction dimension is the input channel (IC) dimension.
+  //   // For Conv2DNhwcHwcfOp, this dimension is already the inner-most
+  //   dimension
+  //   // of the input image, and the penultimate dimension of the kernel --
+  //   // exactly what we want. For Conv2DNchwFchwOp, can the tensor dimensions
+  //   be
+  //   // permuted in DMA to get them in the correct positions? For the image
+  //   // tensor, only if H*W is a nice power of 2 (DMA constraint). For
+  //   kernels,
+  //   // it requires h*w is a nice power of 2 -- unlikely, we typically have
+  //   // h=w=3. The dimension permutations will therefore often therefore need
+  //   to
+  //   // be done on the core. We leave this for future work, the expectation
+  //   for
+  //   // now is that models have been transformed at a high level to avoid
+  //   // channel-first convolutions.
+  //   return linalgOp.emitError(
+  //       "Only channel-last convolution supported currently.");
+  // }
+
+  // // [N, OH, OW, C, KW, HW]
+  else if (isDepthwiseConv2DNhwcHwcOp(linalgOp.getOperation())) {
     // Notes
     // =====
     // A property of depthwise convolution is that it can't be expressed in
@@ -943,6 +950,12 @@ static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
           entryPointFn, genericOp, useLowerToAIEPipeline, targetDevice, numRows,
           numCols, enableAMDAIEUkernels);
     }
+  } else if (isConvOp(genericOp)) {
+    // Current tiling strategy is based on llvm-cpu ConvTileAndDecomposeExpert.
+    if (passPipeline == TilePassPipeline::ConvDecomposePipeline)
+      return setRootConfigForConvDecomposePipeline(
+          entryPointFn, cast<linalg::LinalgOp>(genericOp.getOperation()),
+          targetDevice);
   } else if (isReductionOp(genericOp)) {
     if (passPipeline == TilePassPipeline::GeneralCopyPipeline) {
       return setRootConfigForReductionCopyPipeline(

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Utils/AMDAIEUtils.cpp

@@ -358,6 +358,53 @@ bool isMatmulWithElementwiseConsumer(linalg::LinalgOp linalgOp) {
   return false;
 }
 
+static bool isIteratorTypesOfConvSame(
+    linalg::GenericOp linalgOp,
+    SmallVector<utils::IteratorType> expectedIteratorTypes) {
+  if (!linalgOp) return false;
+  SmallVector<utils::IteratorType> iteratorTypes =
+      linalgOp.getIteratorTypesArray();
+  if (iteratorTypes.size() != expectedIteratorTypes.size()) return false;
+  for (unsigned i = 0, n = expectedIteratorTypes.size(); i < n; i++) {
+    if (iteratorTypes[i] != expectedIteratorTypes[i]) return false;
+  }
+  return true;
+}
+
+bool isConv2DNhwcHwcfOp(Operation *linalgOp) {
+  if (isa<linalg::Conv2DNhwcHwcfOp>(linalgOp)) return true;
+  SmallVector<utils::IteratorType> expectedIteratorTypes = {
+      utils::IteratorType::parallel,  utils::IteratorType::parallel,
+      utils::IteratorType::parallel,  utils::IteratorType::parallel,
+      utils::IteratorType::reduction, utils::IteratorType::reduction,
+      utils::IteratorType::reduction};
+  return isIteratorTypesOfConvSame(dyn_cast<linalg::GenericOp>(linalgOp),
+                                   expectedIteratorTypes);
+}
+
+bool isDepthwiseConv2DNhwcHwcOp(Operation *linalgOp) {
+  if (isa<linalg::DepthwiseConv2DNhwcHwcOp>(linalgOp)) return true;
+  SmallVector<utils::IteratorType> expectedIteratorTypes(
+      cast<ShapedType>(linalgOp->getResult(0).getType()).getRank(),
+      utils::IteratorType::parallel);
+  expectedIteratorTypes.push_back(utils::IteratorType::reduction);
+  expectedIteratorTypes.push_back(utils::IteratorType::reduction);
+  return isIteratorTypesOfConvSame(dyn_cast<linalg::GenericOp>(linalgOp),
+                                   expectedIteratorTypes);
+}
+
+/// Utility to identify whether a linalg op is a broad concept conv op.
+bool isConvOp(linalg::GenericOp linalgOp) {
+  if (!linalgOp) return false;
+  // Test the body of the generic to indeed be what we expect for a matmul.
+  Block *body = linalgOp.getBlock();
+  auto yieldOp = cast<linalg::YieldOp>(body->getTerminator());
+  Value yieldVal = yieldOp.getOperand(0);
+  if (!bodyMatcherForMatmulLikeOps(yieldVal, body)) return false;
+
+  return isConv2DNhwcHwcfOp(linalgOp) || isDepthwiseConv2DNhwcHwcOp(linalgOp);
+}
+
 /// Utility to identify if `linalgOp` is a supported reduction op. Currently,
 /// we are using strict conditions for reduction op matching.
 bool isReductionOp(linalg::LinalgOp linalgOp) {

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Utils/AMDAIEUtils.h

@@ -94,6 +94,10 @@ bool isMatmulWithElementwiseConsumer(linalg::LinalgOp linalgOp);
 /// Utility to identify if `linalgOp` is a supported reduction op.
 bool isReductionOp(linalg::LinalgOp linalgOp);
 
+bool isConv2DNhwcHwcfOp(Operation *linalgOp);
+bool isDepthwiseConv2DNhwcHwcOp(Operation *linalgOp);
+bool isConvOp(linalg::GenericOp linalgOp);
+
 /// Utility to convert a `uint32_t` value into a hex string.
 std::string utohexstr(uint32_t value, size_t width, bool header = true,
                       bool lowercase = false);

runtime/src/iree-amd-aie/driver/xrt-lite/CMakeLists.txt

@@ -40,6 +40,7 @@ iree_cc_library(
     iree::base
     iree::base::core_headers
     iree::hal::utils::deferred_command_buffer
+    iree::hal::utils::queue_emulation
     iree::hal::utils::semaphore_base
     iree::base::internal::flatcc::parsing
     iree-amd-aie::schemas::pdi_executable_def_c_fbs

runtime/src/iree-amd-aie/driver/xrt-lite/device.cc

@@ -15,6 +15,7 @@
 #include "iree-amd-aie/driver/xrt-lite/util.h"
 #include "iree/hal/utils/deferred_command_buffer.h"
 #include "iree/hal/utils/deferred_work_queue.h"
+#include "iree/hal/utils/queue_emulation.h"
 
 #define ARENA_BLOCK_SIZE (32 * 1024)
 
@@ -98,8 +99,9 @@ static iree_status_t iree_hal_xrt_lite_device_create_command_buffer(
 }
 
 static iree_status_t iree_hal_xrt_lite_device_create_semaphore(
-    iree_hal_device_t* base_device, uint64_t initial_value,
-    iree_hal_semaphore_flags_t flags, iree_hal_semaphore_t** out_semaphore) {
+    iree_hal_device_t* base_device, iree_hal_queue_affinity_t queue_affinity,
+    uint64_t initial_value, iree_hal_semaphore_flags_t flags,
+    iree_hal_semaphore_t** out_semaphore) {
   IREE_TRACE_ZONE_BEGIN(z0);
 
   iree_hal_xrt_lite_device* device = IREE_HAL_XRT_LITE_CHECKED_VTABLE_CAST(
@@ -190,8 +192,9 @@ static iree_status_t iree_hal_xrt_lite_device_queue_alloca(
   iree_hal_xrt_lite_device* device = IREE_HAL_XRT_LITE_CHECKED_VTABLE_CAST(
       base_device, iree_hal_xrt_lite_device_vtable, iree_hal_xrt_lite_device);
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
-      z0, iree_hal_semaphore_list_wait(wait_semaphore_list,
-                                       iree_infinite_timeout()));
+      z0,
+      iree_hal_semaphore_list_wait(wait_semaphore_list, iree_infinite_timeout(),
+                                   IREE_HAL_WAIT_FLAG_DEFAULT));
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_hal_allocator_allocate_buffer(device->device_allocator, params,
                                              allocation_size, out_buffer));
@@ -207,8 +210,9 @@ static iree_status_t iree_hal_xrt_lite_device_queue_dealloca(
     const iree_hal_semaphore_list_t wait_semaphore_list,
     const iree_hal_semaphore_list_t signal_semaphore_list,
     iree_hal_buffer_t* buffer, iree_hal_alloca_flags_t flags) {
-  IREE_RETURN_IF_ERROR(iree_hal_semaphore_list_wait(wait_semaphore_list,
-                                                    iree_infinite_timeout()));
+  IREE_RETURN_IF_ERROR(
+      iree_hal_semaphore_list_wait(wait_semaphore_list, iree_infinite_timeout(),
+                                   IREE_HAL_WAIT_FLAG_DEFAULT));
   iree_status_t status = iree_hal_semaphore_list_signal(signal_semaphore_list);
   return status;
 }

runtime/src/iree-amd-aie/driver/xrt-lite/direct_command_buffer.cc

@@ -64,8 +64,11 @@ iree_status_t iree_hal_xrt_lite_direct_command_buffer_create(
   command_buffer->host_allocator = host_allocator;
   command_buffer->device = device;
   iree_arena_initialize(block_pool, &command_buffer->arena);
-  iree_status_t status =
-      iree_hal_resource_set_allocate(block_pool, &command_buffer->resource_set);
+  iree_status_t status = iree_ok_status();
+  if (!iree_all_bits_set(mode, IREE_HAL_COMMAND_BUFFER_MODE_UNRETAINED)) {
+    status = iree_hal_resource_set_allocate(block_pool,
+                                            &command_buffer->resource_set);
+  }
   if (iree_status_is_ok(status)) {
     *out_command_buffer = &command_buffer->base;
   } else {

runtime/src/iree-amd-aie/driver/xrt/CMakeLists.txt

@@ -50,6 +50,7 @@ iree_cc_library(
     iree::hal::utils::deferred_command_buffer
     iree::hal::utils::file_transfer
     iree::hal::utils::files
+    iree::hal::utils::queue_emulation
     iree::hal::utils::semaphore_base
     iree::hal
     iree-amd-aie::schemas::xrt_executable_def_c_fbs