[DNS] Treating clones like transfers during resource usage analysis.

iree-org · Feb 27, 2025 · 54eeac4 · 54eeac4
1 parent 8ba9c20
commit 54eeac4
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
@@ -350,11 +350,58 @@ class ValueResourceUsage : public AbstractResourceUsage<DFX::ValueElement> {
           getState() ^= resultUsage.getState();
         })
         .Case([&](IREE::Stream::AsyncCloneOp op) {
+#if 0
           removeAssumedBits(NOT_TRANSFER_WRITE);
           auto &sourceUsage = solver.getElementFor<ValueResourceUsage>(
               *this, Position::forValue(op.getSource()),
               DFX::Resolution::OPTIONAL);
           getState() ^= sourceUsage.getState();
+#else
+          removeAssumedBits(NOT_TRANSFER_WRITE);
+          auto &sourceUsage = solver.getElementFor<ValueResourceUsage>(
+              *this, Position::forValue(op.getSource()),
+              DFX::Resolution::OPTIONAL);
+          bool isSourceStaging = !(sourceUsage.isAssumed(NOT_STAGING_READ) &&
+                                   sourceUsage.isAssumed(NOT_STAGING_WRITE));
+          bool isTargetStaging =
+              !(isAssumed(NOT_STAGING_READ) && isAssumed(NOT_STAGING_WRITE));
+          if (isSourceStaging != isTargetStaging) {
+            // Can't transition staging across transfers.
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueResourceUsage] skipping transfer source: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          // TODO(benvanik): remove kFavorTransients.
+          bool isSourceExternal = !sourceUsage.isAssumed(NOT_EXTERNAL);
+          bool isTargetInternal = isAssumed(NOT_EXTERNAL);
+          if (kFavorTransients && isSourceExternal && isTargetInternal) {
+            LLVM_DEBUG({
+              llvm::dbgs()
+                  << "[ValueResourceUsage] skipping forward prop of external "
+                     "into internal due to kFavorTransients/device-change: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          auto newState = getState();
+          llvm::errs() << "YYYYY sourceUsage\n";
+          sourceUsage.dump(solver.getAsmState());
+          newState ^= sourceUsage.getState();
+          if (!newState.isValidState()) {
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueResourceUsage] skipping update from "
+                              "producer as it would create an invalid state: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          getState() = newState;
+#endif
         })
         .Case([&](IREE::Stream::AsyncSliceOp op) {
           removeAssumedBits(NOT_TRANSFER_WRITE);
@@ -690,10 +737,54 @@ class ValueResourceUsage : public AbstractResourceUsage<DFX::ValueElement> {
         })
         .Case([&](IREE::Stream::AsyncCloneOp op) {
           removeAssumedBits(NOT_TRANSFER_READ);
+#if 0
+          // token change because bots are dumb
+          // auto &resultUsage = solver.getElementFor<ValueResourceUsage>(
+          //     *this, Position::forValue(op.getResult()),
+          //     DFX::Resolution::OPTIONAL);
+          // getState() ^= resultUsage.getState();
           auto &resultUsage = solver.getElementFor<ValueResourceUsage>(
               *this, Position::forValue(op.getResult()),
               DFX::Resolution::OPTIONAL);
-          getState() ^= resultUsage.getState();
+          bool isSourceStaging =
+              !(isAssumed(NOT_STAGING_READ) && isAssumed(NOT_STAGING_WRITE));
+          bool isTargetStaging = !(resultUsage.isAssumed(NOT_STAGING_READ) &&
+                                   resultUsage.isAssumed(NOT_STAGING_WRITE));
+          if (isSourceStaging != isTargetStaging) {
+            // Can't transition staging across transfers.
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueResourceUsage] skipping clone target: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          // TODO(benvanik): remove kFavorTransients.
+          bool isSourceInternal = isAssumed(NOT_EXTERNAL);
+          bool isTargetExternal = !resultUsage.isAssumed(NOT_EXTERNAL);
+          if (kFavorTransients && isSourceInternal && isTargetExternal) {
+            LLVM_DEBUG({
+              llvm::dbgs()
+                  << "[ValueResourceUsage] skipping back prop of external into "
+                     "internal due to kFavorTransients: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          auto newState = getState();
+          newState ^= resultUsage.getState();
+          if (!newState.isValidState()) {
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueResourceUsage] skipping update from use "
+                              "as it would create an invalid state: ";
+              op.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            return;
+          }
+          getState() = newState;
+#endif
         })
         .Case([&](IREE::Stream::AsyncSliceOp op) {
           removeAssumedBits(NOT_TRANSFER_READ);
@@ -812,6 +903,10 @@ class ValueResourceUsage : public AbstractResourceUsage<DFX::ValueElement> {
         .Case([&](IREE::Stream::AsyncDispatchOp op) {
           removeAssumedBits(NOT_DISPATCH_READ);
           for (auto result : op.getOperandTiedResults(operandIdx)) {
+            llvm::errs() << "XXXX tied result for operand " << operandIdx
+                         << "\n";
+            op.dump();
+            result.dump();
             removeAssumedBits(NOT_MUTATED | NOT_DISPATCH_WRITE);
             auto &resultUsage = solver.getElementFor<ValueResourceUsage>(
                 *this, Position::forValue(result), DFX::Resolution::REQUIRED);
@@ -939,4 +1034,6 @@ LogicalResult ResourceUsageAnalysis::run() {
   return solver.run();
 }
 
+void ResourceUsageAnalysis::print(llvm::raw_ostream &os) { solver.print(os); }
+
 } // namespace mlir::iree_compiler::IREE::Stream
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.h b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.h
@@ -70,6 +70,9 @@ class ResourceUsageAnalysis {
   // May fail if analysis cannot be completed due to unsupported or unknown IR.
   LogicalResult run();
 
+  // Prints the analysis results to |os|.
+  void print(llvm::raw_ostream &os);
+
   // Returns the analyzed resource usage of the |value| resource.
   // May return `Unknown` if analysis failed for the given value due to usage
   // that lead to indeterminate results (such as indirect access).

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp
@@ -468,6 +468,8 @@ struct RefineUsagePass
       return signalPassFailure();
     }
 
+    analysis.print(llvm::errs());
+
     // Query and apply analysis results to all resources in the program.
     RewritePatternSet patterns(&getContext());
     insertUsageRefinementPatterns(&getContext(), analysis, patterns);