Skip to content

Commit

Permalink
[Codegen][GPU] Do not promote fill matmul operands (#18722)
Browse files Browse the repository at this point in the history
If an operand is a fill, there is no reason to promote the operand.
  • Loading branch information
qedawkins authored Oct 9, 2024
1 parent 0e16a89 commit 1b719b3
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,15 @@ void promoteOperand(OpBuilder &builder, Operation *op, unsigned index) {
Value operand = op->getOperand(index);

if (auto producer = operand.getDefiningOp<TilingInterface>()) {
// Skip promotion of fills.
if (isa<linalg::FillOp>(producer)) {
return;
}
if (auto generic = dyn_cast<linalg::GenericOp>(&*producer)) {
if (linalg::isaFillOpInterface(generic)) {
return;
}
}
setLoweringConfig(producer, IREE::GPU::DerivedThreadConfigAttr::get(
builder.getContext()));
return;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,21 @@ func.func @lhs_only_matmul(%a: tensor<32x1024xf32>, %b: tensor<1024x128xf32>) ->
// CHECK-SAME: %[[B:[A-Za-z0-9]+]]: tensor<1024x128xf32>
// CHECK-DAG: %[[PA:.+]] = linalg.copy {{.*}} ins(%[[A]] : tensor<32x1024xf32>)
// CHECK: linalg.generic {{.*}} ins(%[[PA]], %[[B]] : tensor<32x1024xf32>, tensor<1024x128xf32>)

// -----

#lowering_config = #iree_gpu.lowering_config<{promote_operands = [0]}>

func.func @no_promote_fill(%b: tensor<128x128xf32>) -> tensor<4x128xf32> {
%cst = arith.constant 0.000000e+00 : f32
%empty = tensor.empty() : tensor<4x128xf32>
%fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<4x128xf32>) -> tensor<4x128xf32>
%mm = linalg.matmul {lowering_config = #lowering_config}
ins(%fill, %b : tensor<4x128xf32>, tensor<128x128xf32>) outs(%fill : tensor<4x128xf32>) -> tensor<4x128xf32>
return %mm : tensor<4x128xf32>
}

// Verify that fills are not promoted.
// CHECK-LABEL: func.func @no_promote_fill
// CHECK-NOT: iree_gpu.derived_thread_config
// CHECK: return

0 comments on commit 1b719b3

Please sign in to comment.