Add TTNN transform to ensure hoisted ops have operands moved to/from CPU properly

env/CMakeLists.txt

@@ -57,7 +57,7 @@ ExternalProject_Add(
     -DMLIR_LINK_MLIR_DYLIB=OFF
     -DMLIR_BUILD_MLIR_C_DYLIB=ON
     # ======================
-    -DCMAKE_BUILD_TYPE=MinSizeRel
+    -DCMAKE_BUILD_TYPE=Debug
     -DLLVM_ENABLE_ASSERTIONS=ON
     -DMLIR_ENABLE_BINDINGS_PYTHON=ON
     -DCMAKE_C_FLAGS=-D_LIBCPP_HAS_NO_LIBRARY_ALIGNED_ALLOCATION

lib/Dialect/TT/IR/TTOpsTypes.cpp

@@ -14,6 +14,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "ttmlir/Dialect/TT/IR/TT.h"
+#include "ttmlir/Dialect/TTNN/Utils/TransformUtils.h"
 #include "ttmlir/Target/Common/Target.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"

lib/Dialect/TTNN/Transforms/Optimizer.cpp

@@ -146,6 +146,8 @@ class TTNNOptimizer : public impl::TTNNOptimizerBase<TTNNOptimizer> {
 public:
   using impl::TTNNOptimizerBase<TTNNOptimizer>::TTNNOptimizerBase;
   void runOnOperation() final {
+    llvm::outs() << "s TTNNOptimizer::rOO\n";
+
     // Generate legal OP configuration candidates.
     // Perform memory layout analysis.
     // Perform final configuration analysis.
@@ -338,6 +340,7 @@ class TTNNOptimizer : public impl::TTNNOptimizerBase<TTNNOptimizer> {
           func.getContext(), funcType.getInputs(), funcResultTypes);
       func.setType(newFuncType);
     });
+    llvm::outs() << "e TTNNOptimizer::rOO\n";
   }
 
 private:

lib/Dialect/TTNN/Transforms/Passes.cpp

@@ -64,11 +64,16 @@ class TTNNDeallocate : public impl::TTNNDeallocateBase<TTNNDeallocate> {
   }
 
   void runOnOperation() final {
+    llvm::outs() << "s TTNNDeallocate::rOO\n";
     ModuleOp module = getOperation();
     IRRewriter rewriter(&getContext());
 
     module->walk([&](func::FuncOp func) {
-      assert(func.getBody().hasOneBlock());
+      if (func.isDeclaration()) {
+        return;
+      }
+      assert(func.getBody().hasOneBlock() &&
+             "found func that didn't have one block!");
       Liveness liveness(func.getOperation());
       const LivenessBlockInfo *livenessInfo =
           liveness.getLiveness(&func.getBody().front());
@@ -111,6 +116,7 @@ class TTNNDeallocate : public impl::TTNNDeallocateBase<TTNNDeallocate> {
         }
       });
     });
+    llvm::outs() << "e TTNNDeallocate::rOO\n";
   }
 };
 
@@ -122,11 +128,17 @@ class TTNNDecomposeLayouts
       TTNNDecomposeLayouts>::TTNNDecomposeLayoutsBase;
 
   void runOnOperation() final {
+    llvm::outs() << "s TTNNDecomposeLayouts::rOO\n";
+
     ModuleOp module = getOperation();
     IRRewriter rewriter(&getContext());
     llvm::SmallVector<Operation *> opsToReplace;
     module->walk([&](func::FuncOp func) {
-      assert(func.getBody().hasOneBlock());
+      if (func.isDeclaration()) {
+        return;
+      }
+      assert(func.getBody().hasOneBlock() &&
+             "found func that didn't have one block!");
       func->walk([&](Operation *op) {
         if (!isa<ttnn::ToLayoutOp>(op)) {
           return;
@@ -139,6 +151,7 @@ class TTNNDecomposeLayouts
                                       rewriter);
       rewriter.eraseOp(op);
     }
+    llvm::outs() << "e TTNNDecomposeLayouts::rOO\n";
   }
 
 private:
@@ -898,6 +911,8 @@ class TTNNCreateInputGenerators
       TTNNCreateInputGenerators>::TTNNCreateInputGeneratorsBase;
 
   void runOnOperation() final {
+    llvm::outs() << "s TTNNCreateInputGenerators::rOO\n";
+
     ModuleOp module = getOperation();
     IRRewriter rewriter(&getContext());
 
@@ -1066,6 +1081,8 @@ class TTNNCreateInputGenerators
           rewriter.getI32IntegerAttr(0));
       rewriter.create<func::ReturnOp>(mainFuncOp->getLoc(), constantZero);
     }
+
+    llvm::outs() << "e TTNNCreateInputGenerators::rOO\n";
   }
 };
 
@@ -1078,6 +1095,8 @@ class TTNNModifySignaturesForDylib
       TTNNModifySignaturesForDylib>::TTNNModifySignaturesForDylibBase;
 
   void runOnOperation() final {
+    llvm::outs() << "s TTNNModifySignaturesForDylib::rOO\n";
+
     ModuleOp module = getOperation();
     IRRewriter rewriter(&getContext());
 
@@ -1158,6 +1177,7 @@ class TTNNModifySignaturesForDylib
       //
       entryBlock.eraseArguments(1, originalFuncType.getInputs().size());
     }
+    llvm::outs() << "e TTNNModifySignaturesForDylib::rOO\n";
   }
 };
 

lib/Dialect/TTNN/Transforms/TTNNLayout.cpp

@@ -5,12 +5,12 @@
 #include "ttmlir/Dialect/TT/IR/TT.h"
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.h"
 #include "ttmlir/Dialect/TTNN/Transforms/Passes.h"
+#include "ttmlir/Dialect/TTNN/Utils/TransformUtils.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "ttmlir/Dialect/TTNN/Utils/Utils.h"
 
 namespace mlir::tt::ttnn {
 #define GEN_PASS_DEF_TTNNLAYOUT
@@ -51,7 +51,7 @@ inline Location appendInputSuffix(Location loc, int64_t operandIndex) {
 // To layout pass
 //===----------------------------------------------------------------------===//
 
-// Converts tensor types to have a ttnn layout attribute with default values
+// Converts tensor types to have a ttnn layout attribute with deault values
 //
 // Example: tensor<15x10x32xf32> -> tensor<15x10x32xf32, ttnn_layout<...>>
 // where ttnn_layout<...> is constructed with default values
@@ -127,7 +127,12 @@ class TTNNLayoutTensorTypeRewriter : public RewritePattern {
     }
     funcOp.setFunctionType(newType);
 
+    if (funcOp.isDeclaration()) {
+      return true;
+    }
+
     Block &entryBlock = funcOp.getBody().front();
+
     for (unsigned i = 0; i < entryBlock.getNumArguments(); ++i) {
       entryBlock.getArgument(i).setType(inputTypes[i]);
     }
@@ -143,6 +148,10 @@ class TTNNLayoutTensorTypeRewriter : public RewritePattern {
     updated |= convertTypes(op->getOperands(), operands);
     updated |= convertTypes(op->getResults(), results);
     updated |= convertFuncType(op, rewriter);
+    if (updated) {
+      llvm::outs() << "TTNNLayoutTensorTypeRewriter::mAR succeeded on op:\n";
+      op->dump();
+    }
     return updated ? success() : failure();
   }
 
@@ -274,8 +283,12 @@ class TTNNLayoutDPSOperandsRewriter
 
   LogicalResult matchAndRewrite(DestinationStyleOpInterface op,
                                 PatternRewriter &rewriter) const final {
+    llvm::outs() << "s TTNNLayoutDPSOperandsRewriter::mAR\n";
+    op.dump();
+
     // To layout op is a special case, we don't want to rewrite it
     if (mlir::isa<ttir::ToLayoutOp>(op.getOperation())) {
+      llvm::outs() << "e -- f TTNNLayoutDPSOperandsRewriter::mAR\n";
       return failure();
     }
 
@@ -325,11 +338,91 @@ class TTNNLayoutDPSOperandsRewriter
         });
       }
     }
+    llvm::outs() << "e -- " << (modified ? "s" : "f")
+                 << " TTNNLayoutDPSOperandsRewriter::mAR\n";
 
     return modified ? success() : failure();
   }
 };
 
+class TTNNLayoutHoistedFuncCallRewriter
+    : public OpRewritePattern<func::CallOp> {
+public:
+  TTNNLayoutHoistedFuncCallRewriter(MLIRContext *ctx)
+      : OpRewritePattern<func::CallOp>(ctx) {}
+
+  // Match and rewrite the CallOp.
+  LogicalResult matchAndRewrite(func::CallOp callOp,
+                                PatternRewriter &rewriter) const override {
+    llvm::outs() << "s TTNNLayoutHoistedFuncCallRewriter::mAR\n";
+
+    auto device = utils::getOrInsertDevice(rewriter, callOp);
+    llvm::outs() << "0\n";
+
+    llvm::outs() << callOp->getName() << "\n";
+    if (!callOp->hasAttr("hoisted_call")) {
+      llvm::outs() << "e -- f TTNNLayoutHoistedFuncCallRewriter::mAR\n";
+      return failure();
+    }
+
+    llvm::outs() << "1\n";
+
+    // Create a FromDevice operation for each operand.
+    SmallVector<Value, 4> fromDeviceOperands;
+    size_t locIdx = 0;
+    for (auto operand : callOp.getOperands()) {
+      // Insert FromDevice op before the operand and collect the new operands.
+      auto fromDeviceOp = rewriter.create<ttnn::FromDeviceOp>(
+          callOp.getLoc(), operand.getType(), operand);
+      fromDeviceOp.dump();
+      Location newLoc = appendInputSuffix(callOp.getLoc(), locIdx++);
+      std::optional<Value> maybeLayoutOp = createToLayoutOp(
+          rewriter, newLoc, operand, BufferType::SystemMemory,
+          nullptr /* tensorMemoryLayoutAttr */, false /* tiled */);
+      Value hostOpValue = maybeLayoutOp.has_value() ? maybeLayoutOp.value()
+                                                    : fromDeviceOp.getResult();
+      fromDeviceOperands.push_back(hostOpValue);
+    }
+
+    llvm::outs() << "2\n";
+
+    // Create the original CallOp with the new operands (FromDevice'd).
+    auto newCallOp = rewriter.create<func::CallOp>(
+        callOp.getLoc(), callOp.getCallee(), callOp.getResultTypes(),
+        fromDeviceOperands);
+
+    llvm::outs() << "3\n";
+
+    // Now, insert ToDevice ops for the results of the CallOp
+    SmallVector<Value, 4> toDeviceResults;
+    for (auto result : newCallOp.getResults()) {
+      // Insert ToDevice op after the result
+      auto toDeviceOp = rewriter.create<ttnn::ToDeviceOp>(
+          callOp.getLoc(), result.getType(), result, device,
+          ttnn::MemoryConfigAttr{});
+      Location newLoc =
+          appendInputSuffix(callOp.getLoc(), result.getResultNumber() + locIdx);
+      std::optional<Value> maybeLayoutOp = createToLayoutOp(
+          rewriter, newLoc, result, BufferType::SystemMemory,
+          nullptr /* tensorMemoryLayoutAttr */, true /* tiled */);
+      Value deviceResultValue = maybeLayoutOp.has_value()
+                                    ? maybeLayoutOp.value()
+                                    : toDeviceOp.getResult();
+      toDeviceResults.push_back(deviceResultValue);
+    }
+    llvm::outs() << "4\n";
+
+    // Replace the original call with the new ToDevice results.
+    rewriter.replaceOp(callOp, toDeviceResults);
+
+    llvm::outs() << "5\n";
+
+    llvm::outs() << "e -- s TTNNLayoutHoistedFuncCallRewriter::mAR\n";
+
+    return success();
+  }
+};
+
 // Updates the layout of the operands of a func::ReturnOp.
 // The intent is to move the result to host.
 class TTNNLayoutFuncReturnRewriter
@@ -340,6 +433,8 @@ class TTNNLayoutFuncReturnRewriter
 
   LogicalResult matchAndRewrite(mlir::func::ReturnOp op,
                                 PatternRewriter &rewriter) const final {
+    llvm::outs() << "s TTNNLayoutFuncReturnRewriter::mAR\n";
+
     bool modified = false;
     for (OpOperand &operand : op->getOpOperands()) {
       Location newLoc =
@@ -353,6 +448,15 @@ class TTNNLayoutFuncReturnRewriter
         modified = true;
       }
     }
+    llvm::outs() << "e -- ";
+    if (modified) {
+      llvm::outs() << "s";
+      op.dump();
+    } else {
+      llvm::outs() << "f";
+    }
+    llvm::outs() << " TTNNLayoutFuncReturnRewriter::mAR\n";
+
     return modified ? success() : failure();
   }
 
@@ -390,6 +494,9 @@ class TTNNLayout : public impl::TTNNLayoutBase<TTNNLayout> {
       // Takes func::Return op and sets layout which will
       // move it's operands to host
       patterns.add<TTNNLayoutFuncReturnRewriter>(&getContext());
+      // Move operands + results of hoisted funcs to and from device
+      // appropriately
+      patterns.add<TTNNLayoutHoistedFuncCallRewriter>(&getContext());
       FrozenRewritePatternSet patternSet(std::move(patterns));
       GreedyRewriteConfig config = GreedyRewriteConfig();
       config.useTopDownTraversal = true;
@@ -398,6 +505,9 @@ class TTNNLayout : public impl::TTNNLayoutBase<TTNNLayout> {
         signalPassFailure();
         return;
       }
+      llvm::outs() << "IR dump:\n";
+      getOperation()->dump();
+      llvm::outs() << "e TTNNLayout::rOO\n";
     }
   }
 

lib/Dialect/TTNN/Transforms/TTNNWorkarounds.cpp

@@ -398,6 +398,7 @@ class TTNNWorkarounds : public impl::TTNNWorkaroundsBase<TTNNWorkarounds> {
   using impl::TTNNWorkaroundsBase<TTNNWorkarounds>::TTNNWorkaroundsBase;
 
   void runOnOperation() final {
+    llvm::outs() << "s TTNNWorkarounds::rOO\n";
     if (decompositionWorkaroundsEnabled) {
       // Placeholder for workaround decomposition patterns.
       RewritePatternSet patterns(&getContext());
@@ -438,6 +439,7 @@ class TTNNWorkarounds : public impl::TTNNWorkaroundsBase<TTNNWorkarounds> {
         return;
       }
     }
+    llvm::outs() << "e TTNNWorkarounds::rOO\n";
   }
 };
 } // namespace mlir::tt::ttnn