Fix to use OpOperand* in ControlPropagationFn

Signed-off-by: yzhang93 <zhyuhang88@gmail.com>
iree-org · Jul 9, 2024 · a99394c · a99394c
1 parent eab50ac
commit a99394c
Show file tree

Hide file tree

Showing 10 changed files with 14 additions and 95 deletions.
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
diff --git a/.idea/iree.iml b/.idea/iree.iml
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/.idea/workspace.xml b/.idea/workspace.xml
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
@@ -253,7 +253,7 @@ void transform_dialect::ApplyBubbleExpandPatternsOp::populatePatterns(
 void transform_dialect::ApplyBubblePackUnpackPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   linalg::populateDataLayoutPropagationPatterns(
-      patterns, [](Operation *op) { return true; });
+      patterns, [](OpOperand *opOperand) { return true; });
 }
 
 void transform_dialect::ApplyFoldReshapeIntoTensorHalInterfacePatternsOp::

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/PackToIntrinsics.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/PackToIntrinsics.cpp
@@ -97,8 +97,8 @@ void PackToIntrinsicsPass::runOnOperation() {
 
   // Run layout propagation patterns to pull in adjacent un-configured ops.
   RewritePatternSet patterns(context);
-  linalg::ControlPropagationFn control = [](Operation *op) -> bool {
-    return !getLoweringConfig(op);
+  linalg::ControlPropagationFn control = [](OpOperand *opOperand) -> bool {
+    return !getLoweringConfig(opOperand->getOwner());
   };
 
   linalg::populateDataLayoutPropagationPatterns(patterns, control);

diff --git a/compiler/src/iree/compiler/GlobalOptimization/DataLayoutPropagation.cpp b/compiler/src/iree/compiler/GlobalOptimization/DataLayoutPropagation.cpp
@@ -23,14 +23,15 @@ struct DataLayoutPropagationPass
     FunctionOpInterface funcOp = getOperation();
 
     RewritePatternSet patterns(context);
-    linalg::populateDataLayoutPropagationPatterns(patterns, [](Operation *op) {
-      // Currently only bubble up/push down pack/unpack through collapse/expand
-      // shape ops.
-      // TODO(#17734): The propagation through expand_shape ops is broken.
-      // Enable the propagation once we find it useful and the upstream issue is
-      // fixed.
-      return isa<tensor::CollapseShapeOp>(op);
-    });
+    linalg::populateDataLayoutPropagationPatterns(
+        patterns, [](OpOperand *opOperand) {
+          // Currently only bubble up/push down pack/unpack through
+          // collapse/expand shape ops.
+          // TODO(#17734): The propagation through expand_shape ops is broken.
+          // Enable the propagation once we find it useful and the upstream
+          // issue is fixed.
+          return isa<tensor::CollapseShapeOp>(opOperand->get().getDefiningOp());
+        });
     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
       funcOp.emitOpError("folding patterns failed");
       return signalPassFailure();

diff --git a/compiler/src/iree/compiler/GlobalOptimization/GeneralizeLinalgNamedOps.cpp b/compiler/src/iree/compiler/GlobalOptimization/GeneralizeLinalgNamedOps.cpp
@@ -40,7 +40,7 @@ void GeneralizeLinalgNamedOpsPass::runOnOperation() {
                         linalg::DivUnsignedOp, linalg::ElemwiseBinaryOp,
                         linalg::ElemwiseUnaryOp, linalg::ExpOp, linalg::FloorOp,
                         linalg::LogOp, linalg::MapOp, linalg::MaxOp,
-                        linalg::MulOp, linalg::NegfOp, linalg::ReduceOp,
+                        linalg::MulOp, linalg::NegFOp, linalg::ReduceOp,
                         linalg::SubOp, linalg::TransposeOp>(
             linalgOp.getOperation())) {
       namedOpCandidates.push_back(linalgOp);

diff --git a/compiler/src/iree/compiler/Preprocessing/Common/ConvertConvToChannelsLast.cpp b/compiler/src/iree/compiler/Preprocessing/Common/ConvertConvToChannelsLast.cpp
@@ -686,7 +686,7 @@ class ConvertConvToChannelsLastPass
       GreedyRewriteConfig config;
       config.maxIterations = GreedyRewriteConfig::kNoLimit;
       linalg::populateDataLayoutPropagationPatterns(
-          patterns, [](Operation *op) { return true; });
+          patterns, [](OpOperand *opOperand) { return true; });
       if (failed(
               applyPatternsAndFoldGreedily(op, std::move(patterns), config))) {
         return signalPassFailure();