Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir][Linalg]: Optimize linalg generic in transform::PromoteOp to avoid unnecessary copies #68555

Merged
merged 1 commit into from
Oct 14, 2023
Merged

[mlir][Linalg]: Optimize linalg generic in transform::PromoteOp to avoid unnecessary copies #68555

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "mlir/Transforms/FoldUtils.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
Expand Down Expand Up @@ -142,6 +143,8 @@ struct LinalgOpInstancePromotionOptions {
const LinalgPromotionOptions &options);
/// SubViews to promote.
MapVector<int64_t, Value> subViews;
/// Subviews operand numbers to copy in using copyInFn.
llvm::SmallSet<int64_t, 4> operandsNumbersToCopyIn;
/// True if the full view should be used for the promoted buffer.
DenseMap<Value, bool> useFullTileBuffers;

Expand Down Expand Up @@ -174,6 +177,11 @@ LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
Operation *op = opOperand.get().getDefiningOp();
if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) {
subViews[operandNumber] = sv;
// In case of linalg generic, copy in only if subview is used in linalg
// payload.
if (!isa<linalg::GenericOp>(linalgOp) ||
Copy link
Contributor

@chelini chelini Oct 10, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is it tied to a linalg.generic? Can the optimization works for any LinalgOp?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@chelini In case of linalg.generic, it is easy to know if an input/output operand are being used by checking if the operand is being used in the payload before except writing to it (which is what I check payloadUsesValueFromOperand).
If case of any LinalgOp, I am not sure it is triavial since some operations doesn't have a payload (i.e. linalg.copy).

One think we can do it to add an interface for each LinalgOp to implement and check if output operand (init) might be used before assignment.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

payloadUsesValueFromOperand is an interface method of LinalgStructuredInterface, and the copy operation implements LinalgStructuredInterface too. Looking at LinalgNamedStructuredOps.yamlgen.td (you find it in the build dir) CopyOp "derives" from LinalgStructuredBase_Op, which implements LinalgStructuredInterface, thus payloadUsesValueFromOperand is well defined also for CopyOp too.

Copy link
Contributor Author

@AviadCo AviadCo Oct 11, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right about CopyOp has payloadUsesValueFromOperand , but for that example the ins operand needs promotation altough payloadUsesValueFromOperand will return false for it since linalg.copy has not payload.
I believe that we want an inferface that for each LinalgOp return if init operands are being used before writen into.

linalgOp.payloadUsesValueFromOperand(&opOperand))
operandsNumbersToCopyIn.insert(operandNumber);
useFullTileBuffers[sv] = vUseFullTileBuffers[operandNumber];
}
}
Expand Down Expand Up @@ -324,6 +332,8 @@ promoteSubViews(ImplicitLocOpBuilder &b,
auto info = promotionInfoMap.find(v.first);
if (info == promotionInfoMap.end())
continue;
if (options.operandsNumbersToCopyIn.count(v.first) == 0)
continue;
if (failed(options.copyInFn(
b, cast<memref::SubViewOp>(v.second.getDefiningOp()),
info->second.partialLocalView)))
Expand Down
1 change: 1 addition & 0 deletions mlir/test/Dialect/GPU/promotion.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@

// RUN: mlir-opt -allow-unregistered-dialect -pass-pipeline='builtin.module(gpu.module(gpu.func(test-gpu-memory-promotion)))' -split-input-file %s | FileCheck %s

gpu.module @foo {
Expand Down
1 change: 0 additions & 1 deletion mlir/test/Dialect/Linalg/promote.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,6 @@ func.func @linalg_generic_update_all_function_inputs_outputs(%arg0: memref<3x4xf
// CHECK: %[[VAL_62:.*]] = memref.subview %[[VAL_61]][0, 0] {{\[}}%[[VAL_52]], %[[VAL_55]]] [1, 1] : memref<?x?xf32, #gpu.address_space<workgroup>> to memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>
// CHECK: memref.copy %[[VAL_3]], %[[VAL_24]] : memref<4x3xf32, strided<[4, 1]>, 1> to memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>
// CHECK: memref.copy %[[VAL_4]], %[[VAL_43]] : memref<4x3xf32, strided<[4, 1]>, 1> to memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>
// CHECK: memref.copy %[[VAL_5]], %[[VAL_62]] : memref<4x3xf32, strided<[4, 1]>, 1> to memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>
// CHECK: linalg.generic {doc = "", indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"], library_call = ""} ins(%[[VAL_24]], %[[VAL_43]] : memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>, memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>) outs(%[[VAL_62]] : memref<?x?xf32, strided<[?, 1], offset: ?>, #gpu.address_space<workgroup>>) {
// CHECK: ^bb0(%[[VAL_63:.*]]: f32, %[[VAL_64:.*]]: f32, %[[VAL_65:.*]]: f32):
// CHECK: %[[VAL_66:.*]] = arith.addf %[[VAL_63]], %[[VAL_64]] : f32
Expand Down