Improve the semantics of PromoteOperandsOp (#12216)

Make the op also return the promoted operation itself in addition to the
promoted operands. This allows chaining transformations properly without
having to rematch.

compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
index eee80f3..e276361 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp
@@ -664,6 +664,10 @@
   return DiagnosedSilenceableFailure::success();
 }
 
+//===----------------------------------------------------------------------===//
+// PromoteOperandsOp
+//===----------------------------------------------------------------------===//
+
 DiagnosedSilenceableFailure transform_dialect::PromoteOperandsOp::applyToOne(
     Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
@@ -672,6 +676,8 @@
   rewriter.setInsertionPoint(target);
   SmallVector<int64_t> indices = llvm::to_vector(getIndices());
   int64_t numOperands = target->getNumOperands();
+
+  results.push_back(target);
   bufferization::BufferizationOptions options;
   for (int64_t index : indices) {
     if ((index >= 0) && (index < numOperands)) {

compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
index 11c85f4..a13519b 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensionsOps.td
@@ -361,8 +361,11 @@
     This op promotes the specified operands of the provided target handle.
 
     #### Return modes
-    This op returns a handle to an allocTensorOp for each of the provided
-    valid indices.
+    This op consume its target handle and returns a new handle to its target handle
+    as well as an allocTensorOp for each of the provided valid indices.
+
+    If the promotion of any specified operand fails to occur, the op definitely
+    fails.
   }];
 
   let arguments = (ins PDL_Operation:$target,

compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_dialect_promote_operands.mlir

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_dialect_promote_operands.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_dialect_promote_operands.mlir
index ab171ce..0e18f1d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_dialect_promote_operands.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/transform_dialect_promote_operands.mlir
@@ -32,8 +32,11 @@
 
   transform.structured.canonicalized_sequence failures(propagate) {
   ^bb1(%variant_op: !pdl.operation):
-    %matmul = transform.structured.match ops{["linalg.matmul"]} in %variant_op : (!pdl.operation) -> !pdl.operation
-    transform.iree.promote_operands %matmul [0, 1] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
+    %matmul = transform.structured.match ops{["linalg.matmul"]} in %variant_op
+      : (!pdl.operation) -> !pdl.operation
+    %promoted_matmul, %alloc_0, %alloc_1 =
+      transform.iree.promote_operands %matmul [0, 1] 
+        : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
   }
 }