[BufferizeToAllocation] Handle pack operands dependent on other op types #972

jtuyls · 2024-12-09T13:04:45Z

Refactor getOperandsFromDefOp to search for and return pack ops up to a certain depth away from the linalg op's inputs and rename to getPackOperands. This enables more complex cases with intermediate extract_slice operations, see the new test cases being added.

newling · 2024-12-09T16:24:48Z

Looks good, minor suggestions in jtuyls#15

One question -- why is the default depth 2, shouldn't it be 1 because the the bufferization in the current pipeline is always on the closest pack op to the matmul? Or maybe there should just be no default.

jtuyls · 2024-12-09T17:21:04Z

Looks good, minor suggestions in jtuyls#15

One question -- why is the default depth 2, shouldn't it be 1 because the the bufferization in the current pipeline is always on the closest pack op to the matmul? Or maybe there should just be no default.

Yeah, maybe it's better to set default to 1 instead.

yzhang93 · 2024-12-09T18:05:20Z

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

-        clEnumValN(mlir::iree_compiler::AMDAIE::BufferizeOperand::DefOp, "def-op",
-                   "Create new allocations for operands from the def ops of a linalg op.")
-    )}]>
+        clEnumValN(mlir::iree_compiler::AMDAIE::BufferizeOperand::Pack, "pack",


I feel pack may not be a good name for this situation. Maybe we should rename these options to something like "LinalgInput"/"LinalgOutput"/"PackInput"?

Done, renamed it.

Thanks! I see you didn't change the other options. Do you think it's more clear if we change option "Input" to "LinalgInput" and "Output" to "LinalgOutput", etc? This is not critical, so can be addressed in the future if needed.

Yeah, I think that makes sense, but I wanted to avoid convoluting this PR with those changes.

yzhang93 · 2024-12-09T18:13:40Z

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

@@ -1,7 +1,7 @@
 // RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-bufferize-to-allocation{memory-space=2 bufferize-operand=input-output}))' --split-input-file %s | FileCheck %s --check-prefix=INPUT-OUTPUT
 // RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-bufferize-to-allocation{memory-space=2 bufferize-operand=input}))' --split-input-file %s | FileCheck %s --check-prefix=INPUT
 // RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-bufferize-to-allocation{memory-space=2 bufferize-operand=output}))' --split-input-file %s | FileCheck %s --check-prefix=OUTPUT
-// RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-bufferize-to-allocation{memory-space=1 bufferize-operand=def-op}))' --split-input-file %s | FileCheck %s --check-prefix=DEF-OP
+// RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-bufferize-to-allocation{memory-space=1 bufferize-operand=pack pack-depth=2}))' --split-input-file --verify-diagnostics %s | FileCheck %s --check-prefix=PACK


I don't think you need to add --verify-diagnostics?

Yeah, initially, I put all test cases in this file, but some fail for some of the runs commands.

yzhang93 · 2024-12-09T18:19:41Z

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

+    if (!currentOp) {
+      return linalgOp.emitOpError()
+             << "operand #" << input.index() << " only has pack ops to depth "
+             << currentLevel << ", but request is for a depth " << depthLevel
+             << " pack op.";


I don't know if it is easy to add a test for this.

It's here: https://github.com/nod-ai/iree-amd-aie/pull/972/files#diff-08dc5392bd52ce08c6e138c89646fbccbec2f0c782d79323431ea1322919539dR102

yzhang93 · 2024-12-09T18:26:15Z

Looks good, minor suggestions in jtuyls#15
One question -- why is the default depth 2, shouldn't it be 1 because the the bufferization in the current pipeline is always on the closest pack op to the matmul? Or maybe there should just be no default.

Yeah, maybe it's better to set default to 1 instead.

I think we may keep default as 2, as currently PackDepth is only used within getPackOperands function and the existing tests are only with pack-depth=2. And option bufferize-operand=pack pack-depth=1 seems to be a specific case for option bufferize-operand=input. Note the linalg input operand could be other operations than pack ops.

jtuyls · 2024-12-09T20:51:26Z

I think we may keep default as 2, as currently PackDepth is only used within getPackOperands function and the existing tests are only with pack-depth=2. And option bufferize-operand=pack pack-depth=1 seems to be a specific case for option bufferize-operand=input. Note the linalg input operand could be other operations than pack ops.

That's why I initially had it as 2, but I think the default is not necessarily the one we're using, but the one that makes sense from a tranformation's perspective, for which I think 1 is a better default as in principal there are more cases with one level of packs than two.

yzhang93 · 2024-12-09T21:06:12Z

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

-        clEnumValN(mlir::iree_compiler::AMDAIE::BufferizeOperand::DefOp, "def-op",
-                   "Create new allocations for operands from the def ops of a linalg op.")
-    )}]>
+        clEnumValN(mlir::iree_compiler::AMDAIE::BufferizeOperand::PackInput, "pack",


I would change this to pack-input.

jtuyls requested review from MaheshRavishankar, nirvedhmeshram, yzhang93 and Abhishek-Varma as code owners December 9, 2024 13:04

jtuyls force-pushed the bufferize-to-allocation-fix branch from 15aa726 to 3844fc4 Compare December 9, 2024 17:28

yzhang93 reviewed Dec 9, 2024

View reviewed changes

yzhang93 requested changes Dec 9, 2024

View reviewed changes

jtuyls force-pushed the bufferize-to-allocation-fix branch from 3844fc4 to 6f66792 Compare December 9, 2024 20:53

yzhang93 reviewed Dec 9, 2024

View reviewed changes

[BufferizeToAllocation] Handle pack operands dependent on other op types

bcfcb4b

jtuyls force-pushed the bufferize-to-allocation-fix branch from b7cb01d to bcfcb4b Compare December 9, 2024 21:25

Merge branch 'main' into bufferize-to-allocation-fix

71660ae

yzhang93 approved these changes Dec 9, 2024

View reviewed changes

jtuyls merged commit f5ab91e into nod-ai:main Dec 9, 2024
7 checks passed

jtuyls deleted the bufferize-to-allocation-fix branch December 10, 2024 07:44

jtuyls mentioned this pull request Dec 11, 2024

Refactor BufferizeOperand naming for clarity #981

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BufferizeToAllocation] Handle pack operands dependent on other op types #972

[BufferizeToAllocation] Handle pack operands dependent on other op types #972

jtuyls commented Dec 9, 2024

newling commented Dec 9, 2024 •

edited

Loading

jtuyls commented Dec 9, 2024

yzhang93 Dec 9, 2024

jtuyls Dec 9, 2024

yzhang93 Dec 9, 2024

jtuyls Dec 9, 2024

yzhang93 Dec 9, 2024

jtuyls Dec 9, 2024

yzhang93 Dec 9, 2024

jtuyls Dec 9, 2024

yzhang93 commented Dec 9, 2024 •

edited

Loading

jtuyls commented Dec 9, 2024

yzhang93 Dec 9, 2024

jtuyls Dec 9, 2024

[BufferizeToAllocation] Handle pack operands dependent on other op types #972

[BufferizeToAllocation] Handle pack operands dependent on other op types #972

Conversation

jtuyls commented Dec 9, 2024

newling commented Dec 9, 2024 • edited Loading

jtuyls commented Dec 9, 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

yzhang93 commented Dec 9, 2024 • edited Loading

jtuyls commented Dec 9, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

newling commented Dec 9, 2024 •

edited

Loading

yzhang93 commented Dec 9, 2024 •

edited

Loading