squash commit

build_tools/ci/cpu_comparison/run.py

@@ -506,6 +506,9 @@ def aie_vs_baseline(
             if config.return_on_fail:
                 raise RuntimeError("Test failed, exiting.")
 
+    if config.verbose:
+        print(f"Test {name} passed.")
+
 
 def aie_vs_llvm_cpu(
     config,
@@ -607,6 +610,7 @@ def run(self, config):
         for name in [
             "conv2d_nhwc_int32",
             "conv2d_nhwc_bf16",
+            "conv2d_nchw_bf16",
             "conv2d_nhwc_q",
             "depthwise_convolution_i32",
         ]:
@@ -616,7 +620,7 @@ def run(self, config):
                 config,
                 test_files_dir / f"{name}.mlir",
                 tile_pipeline="conv-decompose",
-                lower_to_aie_pipeline="air",
+                lower_to_aie_pipeline="objectFifo",
                 n_repeats=n_conv_repeats,
             )
 
@@ -735,7 +739,7 @@ def run(self, config):
             config,
             test_files_dir / f"conv2d_nhwc_int8.mlir",
             tile_pipeline="conv-decompose",
-            lower_to_aie_pipeline="air",
+            lower_to_aie_pipeline="objectFifo",
             n_repeats=2,
         )
 
@@ -848,6 +852,24 @@ def all_tests(
             error_string += f"\n   {test} ({count} times)."
         raise RuntimeError(error_string)
 
+    else:
+        if verbose:
+            unicorn = r"""
+                    /
+               ,.. /
+             ,'   ';
+  ,,.__    _,' /';  .
+ :','  ~~~~    '. '~ ----- "All tests passed."
+:' (   )         )::,
+'. '. .=----=..-~  .;'
+ '  ;'  ::   ':.  '"
+   (:   ':    ;)
+    \\   '"  ./
+     '"      '"
+        """
+            print(unicorn)
+
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -882,9 +904,9 @@ def all_tests(
         default=0,
         help=dedent(
             """
-            Verbosity level. Currently 
-            0: total silence. 
-            1 (-v) : almost everything. 
+            Verbosity level. Currently
+            0: total silence.
+            1 (-v) : almost everything.
             2 (-vv) : everything.
             """
         ),

build_tools/ci/cpu_comparison/test_files/conv2d_nchw_bf16.mlir

@@ -0,0 +1,11 @@
+// These 2 lines are required by the script which generates input data:
+// input 2x32x14x14xbf16
+// input 64x32x3x3xbf16
+
+func.func @conv_2d_nhwc_hwcf(%arg0: tensor<2x32x14x14xbf16>, %arg1: tensor<64x32x3x3xbf16>) -> tensor<2x64x12x12xf32> {
+  %cst = arith.constant 0.0 : f32
+  %0 = tensor.empty() : tensor<2x64x12x12xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<2x64x12x12xf32>) -> tensor<2x64x12x12xf32>
+  %2 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, %arg1 : tensor<2x32x14x14xbf16>, tensor<64x32x3x3xbf16>) outs(%1 : tensor<2x64x12x12xf32>) -> tensor<2x64x12x12xf32>
+  return %2 : tensor<2x64x12x12xf32>
+}

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

@@ -7,8 +7,11 @@
 #include "iree-amd-aie/IR/AMDAIEAttrs.h"
 #include "iree-amd-aie/Transforms/Passes.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/Pass/Pass.h"
 
 #define DEBUG_TYPE "iree-amdaie-pack-and-transpose"
@@ -20,7 +23,6 @@ namespace {
 static FailureOr<linalg::PackResult> applyPackOnLinalgOp(
     RewriterBase &rewriter, linalg::LinalgOp op,
     SmallVector<OpFoldResult> packedSizes) {
-  // Fail on mismatched number of pack sizes.
   if (packedSizes.size() != op.getNumLoops()) {
     op->emitOpError(
         "requires number of packed sizes match the number of loops (")
@@ -29,12 +31,14 @@ static FailureOr<linalg::PackResult> applyPackOnLinalgOp(
   }
 
   rewriter.setInsertionPoint(op);
-  FailureOr<linalg::PackResult> packResult =
+  FailureOr<linalg::PackResult> maybePackResult =
       linalg::pack(rewriter, op, packedSizes);
-  if (failed(packResult)) {
+  if (failed(maybePackResult)) {
     op->emitOpError("failed to pack the operation");
     return failure();
   }
+
+  linalg::PackResult packResult = maybePackResult.value();
   return packResult;
 }
 
@@ -64,6 +68,10 @@ void AMDAIEPackAndTransposePass::runOnOperation() {
       linalgOp = op;
       return WalkResult::interrupt();
     }
+    if (isa<linalg::ConvolutionOpInterface>(op.getOperation())) {
+      linalgOp = op;
+      return WalkResult::interrupt();
+    }
     return WalkResult::advance();
   });
 
@@ -75,6 +83,7 @@ void AMDAIEPackAndTransposePass::runOnOperation() {
   // Step 1. Before packing the operation, we will prefetch the lowering and
   // packing config.
   auto config = getLoweringConfig<IREE::Codegen::LoweringConfigAttr>(linalgOp);
+
   auto packingConfig = getPackingConfig(linalgOp);
 
   if (!config || !packingConfig) {
@@ -87,6 +96,12 @@ void AMDAIEPackAndTransposePass::runOnOperation() {
   // Extract packing config from the `linalgOp`.
   PackingConfigPackingLevelAttr packCfg =
       packingConfig.getPackingConfigVals(packLevel);
+
+  if (!packCfg) {
+    funcOp->emitOpError("failed to get pack config for pack level ")
+        << packLevel;
+    return signalPassFailure();
+  }
   SmallVector<OpFoldResult> packedSizes =
       getAsIndexOpFoldResult(context, packCfg.getPackedSizes());
 

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

@@ -142,6 +142,36 @@ LogicalResult unPackDmaInputs(IREE::LinalgExt::UnPackOp unPackOp,
   return success();
 }
 
+// Will need funky rebase after
+// https://github.com/nod-ai/iree-amd-aie/pull/776/files
+LogicalResult setDmaInputsFromAllocOp(memref::AllocOp allocOp,
+                                      SmallVector<OpFoldResult> &offsets,
+                                      SmallVector<OpFoldResult> &sizes,
+                                      SmallVector<OpFoldResult> &strides) {
+  auto ctx = allocOp.getContext();
+  auto [stridesI64, baseOffset] = getStridesAndOffset(allocOp.getType());
+  if (baseOffset != 0) {
+    auto message = llvm::formatv(
+        "with non-zero base offset {0} is not supported by the "
+        "current pass, requires testing and possible code changes.",
+        baseOffset);
+    return allocOp->emitOpError(message);
+  }
+  strides = getAsIndexOpFoldResult(ctx, stridesI64);
+  auto sizesI64 = allocOp.getType().getShape();
+  if (llvm::any_of(sizesI64,
+                   [](int64_t size) { return ShapedType::isDynamic(size); })) {
+    return allocOp->emitOpError(
+        "with dynamic shape is not supported by air dma op.");
+  }
+  sizes = getAsIndexOpFoldResult(ctx, sizesI64);
+  // Alloc Op has no offsets.
+  for (int i = 0; i < sizes.size(); i++) {
+    offsets.push_back(getAsIndexOpFoldResult(ctx, 0));
+  }
+  return success();
+}
+
 /// Examines an input/output of a pack/unpack op and provides the
 /// corresponding offsets, sizes and strides required by the dma op
 LogicalResult setDmaInputs(Operation *&operandOp,
@@ -150,28 +180,18 @@ LogicalResult setDmaInputs(Operation *&operandOp,
                            SmallVector<OpFoldResult> &strides) {
   MLIRContext *ctx = operandOp->getContext();
   if (auto allocOp = dyn_cast<memref::AllocOp>(operandOp)) {
-    auto [stridesI64, baseOffset] = getStridesAndOffset(allocOp.getType());
-    if (baseOffset != 0) {
-      auto message = llvm::formatv(
-          "with non-zero base offset {0} is not supported by the "
-          "current pass, requires testing and possible code changes.",
-          baseOffset);
-      return allocOp->emitOpError(message);
-    }
-    strides = getAsIndexOpFoldResult(ctx, stridesI64);
-    auto sizesI64 = allocOp.getType().getShape();
-    if (llvm::any_of(sizesI64, [](int64_t size) {
-          return ShapedType::isDynamic(size);
-        })) {
-      return allocOp->emitOpError(
-          "with dynamic shape is not supported by dma op.");
-    }
-    sizes = getAsIndexOpFoldResult(ctx, sizesI64);
-    // Alloc Op has no offsets.
-    for (int i = 0; i < sizes.size(); i++) {
-      offsets.push_back(getAsIndexOpFoldResult(ctx, 0));
+    return setDmaInputsFromAllocOp(allocOp, offsets, sizes, strides);
+  }
+  if (auto castOp = dyn_cast<memref::CastOp>(operandOp)) {
+    Value castOperand = castOp.getOperand();
+    Operation *castOperandOp = castOperand.getDefiningOp();
+    if (castOperandOp) {
+      memref::AllocOp allocOp = dyn_cast<memref::AllocOp>(castOperandOp);
+      if (allocOp) {
+        operandOp = castOperandOp;
+        return setDmaInputsFromAllocOp(allocOp, offsets, sizes, strides);
+      }
     }
-    return success();
   }
   if (auto subviewOp = dyn_cast<memref::SubViewOp>(operandOp)) {
     auto mixedStrides = subviewOp.getMixedStrides();
@@ -321,6 +341,34 @@ void AMDAIEPackToDmaPass::runOnOperation() {
   MLIRContext *context = &getContext();
   IRRewriter rewriter(context);
 
+  // As a first pass, convert all linalg.copy to iree_linalg_ext.pack.
+  // This is quite simple, the iree_linalg_ext.pack is just an identity pack.
+  // Doing this, we can bootstrap the work already done lowering
+  // the pack op to dmas.
+  auto walkCopyToPackResult =
+      getOperation()->walk([&rewriter](linalg::CopyOp copyOp) -> WalkResult {
+        (void)rewriter;
+        if (copyOp.getNumOperands() != 2 || copyOp.getNumResults() != 0) {
+          copyOp.emitOpError("has ")
+              << copyOp.getNumOperands() << " operands and "
+              << copyOp.getNumResults()
+              << " results. It must have 2 operands and 0 results to convert "
+                 "to an iree.linalg_ext.pack op.";
+          return WalkResult::interrupt();
+        }
+        // Setting up the 'identity' pack:
+        ArrayRef<int64_t> innerDimsPos{};
+        ArrayRef<OpFoldResult> innerTiles{};
+        std::optional<Value> paddingValue{};
+
+        rewriter.setInsertionPoint(copyOp);
+        rewriter.replaceOpWithNewOp<IREE::LinalgExt::PackOp>(
+            copyOp, copyOp.getOperand(0), copyOp.getOperand(1), innerDimsPos,
+            innerTiles, paddingValue);
+        return WalkResult::advance();
+      });
+  if (walkCopyToPackResult.wasInterrupted()) return signalPassFailure();
+
   auto walkResult = getOperation()->walk(
       [&rewriter](IREE::LinalgExt::PackOp op) -> WalkResult {
         if (failed(rewriteAsDma(op, rewriter))) {