[Encoding] Remove ambiguity from encoding propagation interface methods (iree-org#21415)

The `EncodingPropagationAttrInterface` has shared methods for
propagation upwards and downwards, and the methods only take a `Value`
parameter, which is not enough to determine (1) Which operation to
propagate through, or (2) Which direction to propagate. This PR splits
the interface methods into separate methods for upwards and downwards
propagation. The upwards propagation methods now take an `OpResult`, and
the downwards propagation methods take an `OpOperand *`, to indicate for
which operation/operand the propagation is happening for.

There are no test changes, because the propagation interface methods are
already tested in `DispatchCreation/test/propagate_encodings.mlir`.

Signed-off-by: Max Dawkins <[email protected]>
Signed-off-by: keshavvinayak01 <[email protected]>

Author: Max191

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td

@@ -345,19 +345,20 @@ def IREEEncoding_PropagationAttrInterface :
   AttrInterface<"EncodingPropagationAttrInterface"> {
   let cppNamespace = "::mlir::iree_compiler::IREE::Encoding";
   let description = [{
-    Interface used to query new encoding attributes that can be propagated to
-    the operands and result of the operation.
+    Interface used to query new encoding attributes resulting from propagation
+    to the operands and results of operations.
   }];
 
   let methods = [
     InterfaceMethod<
         [{
-          Returns true if the encodings can be propagated across the operation.
+          Returns true if the encoding can be propagated down through the
+          target's owner operation.
         }],
         /*retTy=*/"bool",
-        /*methodName=*/"isPropagable",
+        /*methodName=*/"isPropagableDown",
         /*args=*/(ins
-          "::mlir::Value":$target
+          "::mlir::OpOperand *":$target
         ),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
@@ -366,14 +367,47 @@ def IREEEncoding_PropagationAttrInterface :
     >,
     InterfaceMethod<
         [{
-          Returns the new encodings for operand and result types for the given
+          Returns true if the encoding can be propagated up through the
+          target's owner operation.
+        }],
+        /*retTy=*/"bool",
+        /*methodName=*/"isPropagableUp",
+        /*args=*/(ins
+          "::mlir::OpResult":$target
+        ),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+            return false;
+        }]
+    >,
+    InterfaceMethod<
+        [{
+          Returns the new encodings for operand and result types for the
+          target's owner operation after propagating the encoding down through
+          the operation.
+        }],
+        /*retTy=*/
+        "llvm::FailureOr<::mlir::iree_compiler::IREE::Encoding::PropagationEncoding>",
+        /*methodName=*/"generateSinkingEncodings",
+        /*args=*/(ins
+          "::mlir::OpOperand *":$target
+        ),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return failure();
+      }]
+    >,
+    InterfaceMethod<
+        [{
+          Returns the new encodings for operand and result types for the
+          target's owner operation after propagating the encoding up through the
           operation.
         }],
         /*retTy=*/
         "llvm::FailureOr<::mlir::iree_compiler::IREE::Encoding::PropagationEncoding>",
-        /*methodName=*/"generateEncodings",
+        /*methodName=*/"generateBubblingEncodings",
         /*args=*/(ins
-          "::mlir::Value":$target
+          "::mlir::OpResult":$target
         ),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h

@@ -70,9 +70,9 @@ struct PropagationResult {
   // the propagation.
   SmallVector<Operation *> generatedEncodingOps;
 
-  // The new corresponding result that is created by the propagation. It is
-  // returned to the caller for further transformation or replacement.
-  Value replacement;
+  // The new results created after propagating an encoding through an operation.
+  // It is returned to the caller for further transformation or replacement.
+  SmallVector<Value> replacements;
 };
 
 } // namespace mlir::iree_compiler::IREE::Encoding

compiler/src/iree/compiler/DispatchCreation/PropagateEncodings.cpp

@@ -46,30 +46,30 @@ struct PropagateEncodingsPass
 
 LogicalResult SwapEncodingOpWithTensorCollapseShapeOp::matchAndRewrite(
     IREE::Encoding::SetEncodingOp encodingOp, PatternRewriter &rewriter) const {
-  Value target = encodingOp.getSource();
+  auto collapseOp =
+      encodingOp.getSource().getDefiningOp<tensor::CollapseShapeOp>();
+  if (!collapseOp) {
+    return rewriter.notifyMatchFailure(encodingOp,
+                                       "expected a collapse_shape producer");
+  }
+  auto target = cast<OpResult>(collapseOp.getResult());
   auto propagationAttrInterface =
       dyn_cast<IREE::Encoding::EncodingPropagationAttrInterface>(
           encodingOp.getResultType().getEncoding());
   if (!propagationAttrInterface ||
-      !propagationAttrInterface.isPropagable(target)) {
+      !propagationAttrInterface.isPropagableUp(target)) {
     return rewriter.notifyMatchFailure(
         encodingOp, "the propagation attribute interface isn't defined or the "
                     "target isn't propagable");
   }
   // Get the encoding attributes for the operands and results of the operation.
   FailureOr<IREE::Encoding::PropagationEncoding> propagationEncodings =
-      propagationAttrInterface.generateEncodings(target);
+      propagationAttrInterface.generateBubblingEncodings(target);
   if (failed(propagationEncodings)) {
     return rewriter.notifyMatchFailure(encodingOp,
                                        "not able to determine propagation "
                                        "attributes for operands and results");
   }
-  auto collapseOp =
-      encodingOp.getSource().getDefiningOp<tensor::CollapseShapeOp>();
-  if (!collapseOp) {
-    return rewriter.notifyMatchFailure(encodingOp,
-                                       "expected a collapse_shape producer");
-  }
   if (!IREE::Flow::isNonNullAndOutsideDispatch(encodingOp) ||
       !IREE::Flow::isNonNullAndOutsideDispatch(collapseOp)) {
     return rewriter.notifyMatchFailure(
@@ -91,7 +91,7 @@ LogicalResult SwapEncodingOpWithTensorCollapseShapeOp::matchAndRewrite(
     return rewriter.notifyMatchFailure(
         encodingOp, "not able to propagate encodings and find replacement");
   }
-  rewriter.replaceOp(encodingOp, maybeResult->replacement);
+  rewriter.replaceOp(encodingOp, maybeResult->replacements[0]);
   return success();
 }
 

compiler/src/iree/compiler/ExternalInterfaces/EncodingExternalModels.cpp

@@ -18,9 +18,9 @@ namespace {
 struct ContractionAttrPropagationInterface
     : public IREE::Encoding::EncodingPropagationAttrInterface::ExternalModel<
           ContractionAttrPropagationInterface, IREE::Encoding::MatmulKAttr> {
-  bool isPropagable(Attribute attr, Value target) const {
+  bool isPropagableUp(Attribute attr, OpResult target) const {
     auto encoding = cast<IREE::Encoding::MatmulKAttr>(attr);
-    Operation *attachedToOperation = target.getDefiningOp();
+    Operation *attachedToOperation = target.getOwner();
     if (!attachedToOperation) {
       return false;
     }
@@ -42,11 +42,11 @@ struct ContractionAttrPropagationInterface
   }
 
   FailureOr<IREE::Encoding::PropagationEncoding>
-  generateEncodings(Attribute attr, Value target) const {
+  generateBubblingEncodings(Attribute attr, OpResult target) const {
     auto encoding = cast<IREE::Encoding::MatmulKAttr>(attr);
     return TypeSwitch<Operation *,
                       FailureOr<IREE::Encoding::PropagationEncoding>>(
-               target.getDefiningOp())
+               target.getOwner())
         .Case<tensor::CollapseShapeOp>([&](auto collapseOp) {
           ArrayRef<int32_t> kDims = encoding.getKDims().asArrayRef();
           SmallVector<ReassociationIndices, 4> reassociationMaps =
@@ -98,7 +98,7 @@ struct ContractionOpPropagationInterface
               loc, resultEncodingType, newEncodingOp,
               collapseOp.getReassociationIndices());
           IREE::Encoding::PropagationResult result;
-          result.replacement = newCollapseOp;
+          result.replacements = {newCollapseOp};
           result.generatedEncodingOps.push_back(newEncodingOp.getDefiningOp());
           return result;
         })