[LinalgFunctionOutlining] Create a pass to outline linalg compute ops #862

Abhishek-Varma · 2024-10-29T12:06:14Z

This PR creates a standalone iree-amdaie-linalg-function-outlining pass which will be used to later enable the larger Matmul + Truncf e2e.

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

newling · 2024-10-29T13:51:33Z

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

+  } else if (isElementwise(computeOp)) {
+    computeName = "_elementwise";
+  } else {
+    return failure();


For matmul transpose (etc): It might be nice if this pass, rather than not failing, failed with a message about needing to be implemented.

I've removed failure() altogether and am now simply skipping those compute ops from the walk before invoking this function to outline. Let me know if that sounds okay to you.

It's ok. I think it'd still be good to improve the possible situation in N months time when someone is trying to understand why matmul transpose is running out of PM, whereas matmul is not -- if rather than silently not outlining matmul transpose, you made the pass fail for unconsidered ops, the person might be able to avoid the situation.

I guess it's a trade off between failing too often when you shouldn't and letting issues silently slip through.

newling · 2024-10-29T14:03:49Z

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

+
+/// Utility to outline the linalg compute op.
+static FailureOr<func::FuncOp> outlinedToAFunction(
+    IRRewriter &rewriter, ModuleOp &moduleOp, linalg::LinalgOp &computeOp) {


I think it's unusual to pass ModuleOp/LinalgOp by reference, they're just glamorous pointers. Is there a reason not to pass them by value? Maybe IRRewriter too (although I'm not sure about this one).

I've removed the references from these ops, except rewriter itself because that uses context and I've seen it being referenced in the codebase. If there's a stronger opinion from anyone to have it removed, please let me know and I shall do that.

newling · 2024-10-29T14:09:23Z

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

+  std::string outlinedFuncName =
+      computeOp->getName().stripDialect().str() + computeName + "_outlined";
+  if (auto outlinedFuncOp = dyn_cast_if_present<func::FuncOp>(
+          moduleOp.lookupSymbol(outlinedFuncName))) {


I think you're assuming that if there are 2+ matmuls in the function, then they are identical. Likewise for elementwise operations -- the assumption here is that the unary operation is the same and the input types are the same. While this is true with the current pipeline, it isn't true in general. I thought one reason for using the function outlining approach as opposed to something like #843 was that it would be robust? So I'm wondering if this can be generalized.

Options 1: what happens in MLIR if 2 functions have the exact same signature and body, is there an MLIR pass to detect this and "CSE" them? If so, then we can call the functions '_outlined_0', '_outlined_1' and then let MLIR merge them where possible.

Option 2 (a bit hacky) give the function names more detail -- include M, N, K in function names for matmuls, include unary operation type in function name.

Note: this is not a request for change, I'm ok with this being left for future work :)

Hi @newling - thank you so much for the nudge in this case!

I initially thought that adding just a clear demarcation using isMatmul and isElementwise would be general enough for a particular dispatch, but missed this point in the generalisation.

I've therefore addressed this comment by indeed going ahead with Option 1 as you suggested - involved playing around with IREE's isStructurallyEquivalentTo but I couldn't make that to work, therefore went ahead with upstream's isEquivalentTo.

Please take a look. I've also added corresponding comments in the codebase.

Nice, I didn't know about isEquivalentTo. This seems to do exactly what we want.

I probably would have tried a slightly different design -- I would have created a func.func for every compute operation, and then compared func.funcs for equivalence (using isReqionEquivalentTo) before returning from outlinedToAFunc. This would require keeping a DenseSet<FuncOp> instead of the DenseMap, which would mean we could erase compute ops immediately rather than waiting till the end of the pass (smaller mental stack while reading code!)

This is just a minor design preference though -- I will make a PR with this change after this lands.

newling · 2024-10-29T14:20:24Z

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 AMDAIELinalgFunctionOutlining :
+  Pass<"iree-amdaie-linalg-function-outlining", "ModuleOp"> {
+  let summary = "Outlining of linalg compute ops";


Can you add some additional information in a description = [{ }] block about which linalg ops specifically are outlined (i.e. linalg.fill is not) and what the motivation for this pass is, and what some assumptions are? Sorry, I know that most of our passes in this file don't contain descriptions, but I think it'd be a good thing if we all started adding high-level information here for people without context.

I have added that. Please check if the wordings are correct.

Thanks.

Is

One assumption this pass currently makes: All elementwise/matmul linalg ops within a dispatch have same body content.

still true? Otherwise looks good.

Same question. Doesnt operation equivalence check the region as well?

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

yzhang93

Some of my comments are still in the other PR that need to be addressed.

newling · 2024-10-30T15:39:20Z

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

+
+namespace {
+
+unsigned uniqueOutlinedMatmul = 0;


Design nit: I would make these (private) members of the Pass, and make outlinedToAFunction a member function of the Pass too. Just because I'm a bit unsure about these non-static global variables, not seen this C++ design before.

newling · 2024-10-30T15:43:25Z

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

+    IRRewriter &rewriter, ModuleOp moduleOp, linalg::LinalgOp computeOp,
+    std::string outlineFuncName,
+    DenseMap<Operation *, std::string> &computeOpToOutlinedFuncMap) {
+  // Check if the compute op is equivalent to a previously outlined compute op.


At this point we expect moduleOp's symbol table to not contain outlineFuncName. I think this is worth asserting, to avoid to unlikely case that there is already a function with the name you've just created. I suppose you could increment the uniqueOutlined index in a while loop until an unused name is found.

newling

Functionally looks good to me. My only remaining comments are mostly minor design things which I'm happy to bikeshed, so accepting as is. But @Abhishek-Varma please ping me if you'd like me to take another look.

MaheshRavishankar · 2024-10-31T03:53:12Z

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

+
+namespace {
+
+unsigned uniqueOutlinedMatmul = 0;


Global values like this arent a good idea. This can cause race conditions when run in multi-threaded context.

I think LLVM already has a way to deduplicate symbols. I dont know how to call it though.

MaheshRavishankar · 2024-10-31T03:56:35Z

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 AMDAIELinalgFunctionOutlining :
+  Pass<"iree-amdaie-linalg-function-outlining", "ModuleOp"> {
+  let summary = "Outlining of linalg compute ops";


Same question. Doesnt operation equivalence check the region as well?

yzhang93 · 2024-10-31T17:17:06Z

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

+    if (!OperationEquivalence::isEquivalentTo(
+            computeOp.getOperation(), op,
+            OperationEquivalence::ignoreValueEquivalence, /*flags=*/nullptr,
+            OperationEquivalence::IgnoreLocations))
+      continue;
+    outlineFuncName = funcName;
+    break;


Why don't you directly check if the isEquivalentTo condition holds, and then update outlineFuncName and break?

yzhang93 · 2024-10-31T17:24:09Z

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

+  DenseMap<Operation *, std::string> computeOpToOutlinedFuncMap;
+  SmallVector<Operation *> toBeErased;
+  moduleOp.walk([&](linalg::LinalgOp computeOp) {
+    // Form outlined FuncName.


Also add comments of what linalg ops are currently outlined.

yzhang93 · 2024-10-31T17:29:08Z

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

+  let summary = "Outlining of linalg compute ops";
+  let description = [{
+    Outlines matmul/elementwise linalg compute ops only. This pass essentially minimises
+    the code footprint overall as the unrolling in the later passes lead to lot repeated


Nit: 1) "loop unrolling"; 2) "lead to lots of repeated codes".

yzhang93 · 2024-10-31T17:34:19Z

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

+  %7 = amdaie.core(%tile, in : [], out : []) {
+    linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]
+    } ins(%A : memref<4xf32>)
+      outs(%C : memref<4xbf16>) {
+    ^bb0(%in: f32, %out: bf16):
+      %8 = arith.truncf %in : f32 to bf16
+      linalg.yield %8 : bf16
+    }
+    amdaie.end
+  }


I think the above three elementwise cores are good enough for test purpose. Why don't you delete this or change this to a matmul.

Abhishek-Varma added 6 commits October 29, 2024 11:56

[WIP] Create a function outlining pass

493e08d

Subsume the compute op and form the function body

9e686fd

After insert-cores fix

1f04f8e

Update lower-to-aie for outlined function

48cbd72

Add lit test

daeb581

Review comments for FuncOutlining

436a402

Abhishek-Varma requested review from MaheshRavishankar, nirvedhmeshram, yzhang93 and jtuyls as code owners October 29, 2024 12:06

Abhishek-Varma mentioned this pull request Oct 29, 2024

[Vectorization][ObjectFifo] Enable larger Matmul + Truncf #856

Open

newling reviewed Oct 29, 2024

View reviewed changes

yzhang93 requested changes Oct 29, 2024

View reviewed changes

Review comments

9bf1eb3

Abhishek-Varma requested review from newling and yzhang93 October 30, 2024 08:27

newling reviewed Oct 30, 2024

View reviewed changes

newling approved these changes Oct 30, 2024

View reviewed changes

MaheshRavishankar requested changes Oct 31, 2024

View reviewed changes

yzhang93 requested changes Oct 31, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[LinalgFunctionOutlining] Create a pass to outline linalg compute ops #862

[LinalgFunctionOutlining] Create a pass to outline linalg compute ops #862

Abhishek-Varma commented Oct 29, 2024

newling Oct 29, 2024

Abhishek-Varma Oct 30, 2024

newling Oct 30, 2024

newling Oct 29, 2024

Abhishek-Varma Oct 30, 2024

newling Oct 29, 2024

Abhishek-Varma Oct 30, 2024

newling Oct 30, 2024

newling Oct 29, 2024

Abhishek-Varma Oct 30, 2024

newling Oct 30, 2024

MaheshRavishankar Oct 31, 2024

yzhang93 left a comment

newling Oct 30, 2024

newling Oct 30, 2024

newling left a comment

MaheshRavishankar Oct 31, 2024

MaheshRavishankar Oct 31, 2024

yzhang93 Oct 31, 2024

yzhang93 Oct 31, 2024

yzhang93 Oct 31, 2024

yzhang93 Oct 31, 2024


		namespace {

		unsigned uniqueOutlinedMatmul = 0;


		namespace {

		unsigned uniqueOutlinedMatmul = 0;

[LinalgFunctionOutlining] Create a pass to outline linalg compute ops #862

Are you sure you want to change the base?

[LinalgFunctionOutlining] Create a pass to outline linalg compute ops #862

Conversation

Abhishek-Varma commented Oct 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhang93 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

newling left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment