Add function outlining pass + e2e larger Matmul + Truncf

build_tools/ci/cpu_comparison/run.py

@@ -783,6 +783,29 @@ def run(self, config):
             output_type=get_output_type(test_name),
         )
 
+        # Large shape Matmul + Truncf
+        generate_matmul_test(test_name, template_name, 128, 128, 256, "bf16", "f32")
+        identity_mat = np.eye(128, dtype=np.float32)
+        ones = np.ones(128 * 128, dtype=np.float32).reshape([128, 128])
+        lhs = ones * 101
+        rhs = identity_mat * 3
+        input_args = generate_inputs(test_name, output_dir, 1, {1: lhs, 2: rhs})
+        aie_vs_baseline(
+            config,
+            test_name,
+            input_args,
+            ones * 302,  # exected output
+            use_ukernel=False,
+            tile_pipeline="pack-peel",
+            lower_to_aie_pipeline="objectFifo",
+            function_name=None,
+            seed=1,
+            rtol=0,
+            atol=0,
+            n_repeats=1,
+            output_type=get_output_type(test_name),
+        )
+
 
 class SmokeSet(TestSet):
     def __init__(self):

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

@@ -0,0 +1,112 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree-amd-aie/IR/AMDAIEOps.h"
+#include "iree-amd-aie/Transforms/AMDAIEUtils.h"
+#include "iree-amd-aie/Transforms/Passes.h"
+#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
+#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+
+#define DEBUG_TYPE "iree-amdaie-function-outlining"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+namespace {
+
+class AMDAIEFunctionOutliningPass
+    : public impl::AMDAIEFunctionOutliningBase<AMDAIEFunctionOutliningPass> {
+ public:
+  AMDAIEFunctionOutliningPass() = default;
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AMDAIEDialect, linalg::LinalgDialect>();
+  }
+
+  void runOnOperation() override;
+};
+
+void AMDAIEFunctionOutliningPass::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+  MLIRContext *context = &getContext();
+  IRRewriter rewriter(context);
+
+  auto outlinedToAFunction = [&](linalg::LinalgOp &computeOp) -> func::FuncOp {
+    // Form outlined FuncName.
+    std::string computeName = "";
+    if (isMatmul(computeOp)) {
+      computeName = "_matmul";
+    } else {
+      // TODO(avarma): Make this better/general.
+      computeName = "_elementwise";
+    }
+    std::string outlinedFuncName =
+        computeOp->getName().stripDialect().str() + computeName + "_outlined";
+    if (auto outlinedFuncOp = dyn_cast_if_present<func::FuncOp>(
+            moduleOp.lookupSymbol(outlinedFuncName)))
+      return outlinedFuncOp;
+
+    // Form outlined FunctionType.
+    SmallVector<Type> inputTypes = llvm::map_to_vector(
+        computeOp.getDpsInputs(), [](Value v) { return v.getType(); });
+    for (Value val : computeOp.getDpsInits())
+      inputTypes.push_back(val.getType());
+    auto outlinedFuncType =
+        FunctionType::get(rewriter.getContext(), inputTypes, {});
+
+    // Form outlined FuncSignature
+    rewriter.setInsertionPointToStart(moduleOp.getBody());
+    auto outlinedFunc = rewriter.create<func::FuncOp>(
+        moduleOp.getLoc(), outlinedFuncName, outlinedFuncType);
+    outlinedFunc.setPrivate();
+
+    // Create an entry func block and map the original operands of the compute
+    // op to the block arguments.
+    Block *outlinedFuncBody = outlinedFunc.addEntryBlock();
+    rewriter.setInsertionPointToStart(outlinedFuncBody);
+    SmallVector<BlockArgument> outlinedFuncArgs =
+        llvm::map_to_vector(outlinedFunc.getArguments(),
+                            [&](BlockArgument bbArg) { return bbArg; });
+    unsigned bbArgIndex = 0;
+    IRMapping operandMap;
+    for (Value origOperand : computeOp.getDpsInputs()) {
+      operandMap.map(origOperand, outlinedFuncArgs[bbArgIndex++]);
+    }
+    for (Value origOperand : computeOp.getDpsInits()) {
+      operandMap.map(origOperand, outlinedFuncArgs[bbArgIndex++]);
+    }
+
+    // Clone the compute op while mapping the operand to the function block
+    // arguments.
+    Operation *clonedComputeOp = rewriter.clone(*computeOp, operandMap);
+
+    // Create terminator op returning the cloned compute op's results.
+    rewriter.setInsertionPointToEnd(outlinedFuncBody);
+    rewriter.create<func::ReturnOp>(clonedComputeOp->getLoc(), ValueRange({}));
+
+    return outlinedFunc;
+  };
+
+  moduleOp.walk([&](linalg::LinalgOp computeOp) {
+    if (isa<linalg::FillOp, linalg::CopyOp>(computeOp))
+      return WalkResult::skip();
+    func::FuncOp outlinedFuncOp = outlinedToAFunction(computeOp);
+    rewriter.setInsertionPoint(computeOp);
+    rewriter.create<func::CallOp>(computeOp.getLoc(), outlinedFuncOp,
+                                  computeOp->getOperands());
+    rewriter.eraseOp(computeOp);
+    return WalkResult::advance();
+  });
+}
+
+}  // namespace
+
+std::unique_ptr<Pass> createAMDAIEFunctionOutliningPass() {
+  return std::make_unique<AMDAIEFunctionOutliningPass>();
+}
+}  // namespace mlir::iree_compiler::AMDAIE

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

@@ -297,6 +297,7 @@ LogicalResult AIEDeviceBuilder::coreFuncCallOpToAIE(
     SymbolTable::setSymbolVisibility(newFnDecl,
                                      SymbolTable::Visibility::Private);
     newFnDecl->setAttr("llvm.bareptr", rewriter.getBoolAttr(true));
+    fnDecl.getBody().cloneInto(&(newFnDecl.getBody()), mapper);
     mapper.map(fnDecl.getOperation(), newFnDecl.getOperation());
     fnDecl = newFnDecl;
   }

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/CMakeLists.txt

@@ -70,6 +70,7 @@ iree_cc_library(
     "AMDAIEDmaToCircularDma.cpp"
     "AMDAIEDmaUtils.cpp"
     "AMDAIEFlattenLogicalObjectFifo.cpp"
+    "AMDAIEFunctionOutlining.cpp"
     "AMDAIEFuseConsumerIntoLoop.cpp"
     "AMDAIEFuseFillIntoForall.cpp"
     "AMDAIEFusePackIntoLoop.cpp"

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

@@ -47,6 +47,7 @@ namespace mlir::iree_compiler::AMDAIE {
 #define GEN_PASS_DEF_AMDAIEDMALOOPSUBSUMPTION
 #define GEN_PASS_DEF_AMDAIEDMATOCIRCULARDMA
 #define GEN_PASS_DEF_AMDAIEFLATTENLOGICALOBJECTFIFO
+#define GEN_PASS_DEF_AMDAIEFUNCTIONOUTLINING
 #define GEN_PASS_DEF_AMDAIEFUSECONSUMERINTOLOOP
 #define GEN_PASS_DEF_AMDAIEFUSEFILLINTOFORALL
 #define GEN_PASS_DEF_AMDAIEFUSEPACKINTOLOOP

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

@@ -560,6 +560,8 @@ void addAMDAIEObjectFifoLoweringPasses(OpPassManager &passManager,
   passManager.addPass(createAMDAIENormalizeLoopBoundsPass());
   passManager.addPass(createAMDAIEInsertCoresPass());
 
+  passManager.addPass(createAMDAIEFunctionOutliningPass());
+
   {
     // Vectorization passes
     OpPassManager &funcPassManager = passManager.nest<func::FuncOp>();

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

@@ -163,6 +163,9 @@ std::unique_ptr<Pass> createAMDAIEDmaToCircularDmaPass();
 /// Create a pass to flatten the logical objectFifos.
 std::unique_ptr<Pass> createAMDAIEFlattenLogicalObjectFifoPass();
 
+/// Create a pass for function outlining.
+std::unique_ptr<Pass> createAMDAIEFunctionOutliningPass();
+
 /// Create a pass to fuse the consumer op into the innermost last scf loop.
 std::unique_ptr<Pass> createAMDAIEFuseConsumerIntoLoopPass(
     AMDAIEFuseConsumerIntoLoopOptions options = {});

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td

@@ -243,6 +243,12 @@ def AMDAIEFlattenLogicalObjectFifo :
   let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIEFlattenLogicalObjectFifoPass()";
 }
 
+def AMDAIEFunctionOutlining :
+  Pass<"iree-amdaie-function-outlining", "ModuleOp"> {
+  let summary = "Function outlining";
+  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIEFunctionOutliningPass()";
+}
+
 def AMDAIEFuseConsumerIntoLoop :
     InterfacePass<"iree-amdaie-fuse-consumer-into-loop", "mlir::FunctionOpInterface"> {
   let summary = "Fuse the consumer operation into the innermost last scf loop.";

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/CMakeLists.txt

@@ -35,6 +35,7 @@ iree_lit_test_suite(
     "dma_loop_subsumption.mlir"
     "dma_to_circular_dma.mlir"
     "flatten_logical_objectfifo.mlir"
+    "function_outlining.mlir"
     "fuse_consumer_into_loop_scf_for.mlir"
     "fuse_consumer_into_loop_scf_forall.mlir"
     "fuse_fill_into_forall.mlir"

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/function_outlining.mlir

@@ -0,0 +1,91 @@
+// RUN: iree-opt --split-input-file --iree-amdaie-function-outlining --verify-diagnostics --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func.func private @generic_matmul_outlined
+// CHECK-SAME:  (%[[LHS:.*]]: memref<1x1x4x8x4x8xbf16>,
+// CHECK-SAME:   %[[RHS:.*]]: memref<1x1x8x4x8x4xbf16>,
+// CHECK-SAME:   %[[OUT:.*]]: memref<1x1x8x8x4x4xf32>) {
+// CHECK:           linalg.generic
+// CHECK-SAME:          ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME:          outs(%[[OUT]] :
+// CHECK:           return
+// CHECK:        }
+// CHECK-LABEL:  func.func @matmul_example
+// CHECK-SAME:   (%[[A:.*]]: memref<1x1x4x8x4x8xbf16>,
+// CHECK-SAME:    %[[B:.*]]: memref<1x1x8x4x8x4xbf16>,
+// CHECK-SAME:    %[[C:.*]]: memref<1x1x8x8x4x4xf32>) {
+// CHECK:            amdaie.core
+// CHECK:               func.call @generic_matmul_outlined(%[[A]], %[[B]], %[[C]])
+// CHECK-NOT:           linalg.generic
+// CHECK:               amdaie.end
+// CHECK:            }
+// CHECK:            return
+// CHECK:        }
+func.func @matmul_example(%A: memref<1x1x4x8x4x8xbf16>, %B: memref<1x1x8x4x8x4xbf16>, %C: memref<1x1x8x8x4x4xf32>) {
+  %c2 = arith.constant 2 : index
+  %c1 = arith.constant 1 : index
+  %tile = amdaie.tile(%c1, %c2)
+  %0 = amdaie.core(%tile, in : [], out : []) {
+    linalg.generic {
+      indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d2, d5, d3, d6, d8)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d2, d1, d4, d5, d8, d7)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d4, d3, d6, d7)>
+                      ],
+      iterator_types = ["parallel", "parallel", "reduction",
+                        "parallel", "parallel", "reduction",
+                        "parallel", "parallel", "reduction"
+                       ]
+    } ins(%A, %B : memref<1x1x4x8x4x8xbf16>, memref<1x1x8x4x8x4xbf16>)
+      outs(%C : memref<1x1x8x8x4x4xf32>) {
+    ^bb0(%in: bf16, %in_17: bf16, %out: f32):
+      %1 = arith.extf %in : bf16 to f32
+      %2 = arith.extf %in_17 : bf16 to f32
+      %3 = arith.mulf %1, %2 : f32
+      %4 = arith.addf %out, %3 : f32
+      linalg.yield %4  : f32
+    }
+    amdaie.end
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @generic_elementwise_outlined
+// CHECK-SAME:  (%[[INPUT:.*]]: memref<1x1x8x8x4x4xf32>,
+// CHECK-SAME:   %[[OUTPUT:.*]]: memref<1x1x8x8x4x4xbf16>) {
+// CHECK:           linalg.generic
+// CHECK-SAME:          ins(%[[INPUT]] :
+// CHECK-SAME:          outs(%[[OUTPUT]] :
+// CHECK:           return
+// CHECK:        }
+// CHECK-LABEL:  func.func @elemwise_example
+// CHECK-SAME:   (%[[A:.*]]: memref<1x1x8x8x4x4xf32>,
+// CHECK-SAME:    %[[C:.*]]: memref<1x1x8x8x4x4xbf16>) {
+// CHECK:            amdaie.core
+// CHECK:               func.call @generic_elementwise_outlined(%[[A]], %[[C]])
+// CHECK-NOT:           linalg.generic
+// CHECK:               amdaie.end
+// CHECK:            }
+// CHECK:            return
+// CHECK:        }
+func.func @elemwise_example(%A: memref<1x1x8x8x4x4xf32>, %C: memref<1x1x8x8x4x4xbf16>) {
+  %c2 = arith.constant 2 : index
+  %c1 = arith.constant 1 : index
+  %tile = amdaie.tile(%c1, %c2)
+  %0 = amdaie.core(%tile, in : [], out : []) {
+    linalg.generic {
+      indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>],
+      iterator_types = ["parallel", "parallel", "parallel",
+                        "parallel", "parallel", "parallel"
+                       ]
+    } ins(%A : memref<1x1x8x8x4x4xf32>)
+      outs(%C : memref<1x1x8x8x4x4xbf16>) {
+    ^bb0(%in: f32, %out: bf16):
+      %1 = arith.truncf %in : f32 to bf16
+      linalg.yield %1 : bf16
+    }
+    amdaie.end
+  }
+  return
+}