support simple strctured ops in layout propagation

yifeizh2 · yifeizh2 · commit 58f363edfb6e · 2024-05-27T01:33:54.000-07:00
diff --git a/include/gc/Analysis/GlobalAnalysis.h b/include/gc/Analysis/GlobalAnalysis.h
@@ -23,6 +23,7 @@
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
 
@@ -34,7 +35,7 @@ using namespace mlir;
 class TensorLayout {
 public:
   TensorLayout(ArrayRef<int64_t> outerAxis, ArrayRef<int64_t> innerAxis,
-               ArrayRef<int64_t> tileSizes) {
+               ArrayRef<OpFoldResult> tileSizes) {
     assert(innerAxis.size() == tileSizes.size());
     for (auto oa : outerAxis) {
       OuterAxis.push_back(oa);
@@ -59,7 +60,7 @@ class TensorLayout {
     SmallVector<int64_t> outerAxis(rank, 0);
     std::iota(outerAxis.begin(), outerAxis.end(), 0);
     return TensorLayout(outerAxis, SmallVector<int64_t>{},
-                        SmallVector<int64_t>{});
+                        SmallVector<OpFoldResult>{});
   }
 
   size_t getTensorRank() const { return OuterAxis.size(); }
@@ -68,7 +69,7 @@ class TensorLayout {
 
   SmallVector<int64_t> getInnerAxis() const { return InnerAxis; }
 
-  SmallVector<int64_t> getTileSizes() const { return TileSizes; }
+  SmallVector<OpFoldResult> getTileSizes() const { return TileSizes; }
 
   friend std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout);
 
@@ -77,7 +78,7 @@ class TensorLayout {
 private:
   SmallVector<int64_t> OuterAxis;
   SmallVector<int64_t> InnerAxis;
-  SmallVector<int64_t> TileSizes;
+  SmallVector<OpFoldResult> TileSizes;
 };
 
 class OperatorLayout {
diff --git a/lib/gc/Analysis/GlobalAnalysis.cpp b/lib/gc/Analysis/GlobalAnalysis.cpp
@@ -16,7 +16,7 @@ namespace gc {
 std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout) {
   SmallVector<int64_t> outerAxis = layout.getOuterAxis();
   SmallVector<int64_t> innerAxis = layout.getInnerAxis();
-  SmallVector<int64_t> tileSizes = layout.getTileSizes();
+  SmallVector<OpFoldResult> tileSizes = layout.getTileSizes();
   ss << "[";
   for (size_t i = 0; i < outerAxis.size(); ++i) {
     if (i != 0) {
@@ -35,7 +35,9 @@ std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout) {
       if (i != 0) {
         ss << ", ";
       }
-      ss << tileSizes[i];
+      if (getConstantIntValue(tileSizes[i]).has_value()) {
+        ss << *getConstantIntValue(tileSizes[i]);
+      }
     }
     ss << "}";
   }
@@ -120,10 +122,10 @@ inferTargetLayout(TensorLayout layoutBase,
   int64_t dimDifference = indexMap.size() - layoutBase.getTensorRank();
   SmallVector<int64_t> baseOuterAxis = layoutBase.getOuterAxis();
   SmallVector<int64_t> baseInnerAxis = layoutBase.getInnerAxis();
-  SmallVector<int64_t> baseTileSizes = layoutBase.getTileSizes();
+  SmallVector<OpFoldResult> baseTileSizes = layoutBase.getTileSizes();
   SmallVector<int64_t> targetOuterAxis;
   SmallVector<int64_t> targetInnerAxis;
-  SmallVector<int64_t> targetTileSizes;
+  SmallVector<OpFoldResult> targetTileSizes;
   DenseMap<int64_t, int64_t> reverseIndexMap =
       getReversedIndexMap(indexMap, layoutBase.getTensorRank());
   for (auto oa : baseOuterAxis) {
@@ -184,18 +186,28 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
       }
 
       // ------ Get Current Op's Suggested Layout & Do Propagation ------
+      IRRewriter rewriter(linalgOp);
       if (mlir::linalg::isaContractionOpInterface(linalgOp)) {
         // query the cost model
         // OperatorLayout suggestedLayout = costModel->queryLayout(linalgOp,
         // curInputLayouts);
 
         // hardcode one for now
         // A side layout, [0, 1, 0, 1]; {32, 32}
-        TensorLayout A_layout({0, 1}, {0, 1}, {32, 32});
+        TensorLayout A_layout(
+            {0, 1}, {0, 1},
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
+                                      rewriter.getIndexAttr(32)});
         // B side layout, [1, 0, 0, 1]; {32, 32}
-        TensorLayout B_layout({1, 0}, {0, 1}, {32, 32});
+        TensorLayout B_layout(
+            {1, 0}, {0, 1},
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
+                                      rewriter.getIndexAttr(32)});
         // C side layout, [0, 1, 0, 1]; {32, 32}
-        TensorLayout C_layout({0, 1}, {0, 1}, {32, 32});
+        TensorLayout C_layout(
+            {0, 1}, {0, 1},
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
+                                      rewriter.getIndexAttr(32)});
         OperatorLayout suggestedLayout({A_layout, B_layout}, {C_layout});
         layout[linalgOp] = suggestedLayout;
       } else {
diff --git a/lib/gc/Transforms/PropagateLayout.cpp b/lib/gc/Transforms/PropagateLayout.cpp
@@ -13,6 +13,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
@@ -30,6 +31,130 @@ using namespace mlir;
 using namespace mlir::arith;
 using namespace mlir::tensor;
 
+static FailureOr<linalg::PackResult> packNamedOp(RewriterBase &rewriter,
+                                                 linalg::LinalgOp linalgOp,
+                                                 OperatorLayout opLayout) {
+  std::cout << "----------------------------------" << std::endl;
+  std::cout << " Visiting op in packNamedOp ";
+  linalgOp->getName().print(llvm::errs());
+  std::cout << std::endl;
+  std::cout << "----------------------------------" << std::endl;
+  Location loc = linalgOp->getLoc();
+  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+  SmallVector<utils::IteratorType> iteratorTypes =
+      linalgOp.getIteratorTypesArray();
+
+  SmallVector<tensor::PackOp> packOps;
+  SmallVector<tensor::UnPackOp> unPackOps;
+  SmallVector<Value> inputsAndInits, results;
+  SmallVector<OpOperand *> initOperands = llvm::to_vector(llvm::map_range(
+      linalgOp.getDpsInitsMutable(), [](OpOperand &o) { return &o; }));
+  SmallVector<OpOperand *> inputOperands = linalgOp.getDpsInputOperands();
+  std::cout << "Num of input operands: " << inputOperands.size() << std::endl;
+  std::cout << "Num of init operands: " << initOperands.size() << std::endl;
+  SmallVector<TensorLayout> inputLayouts = opLayout.getSupportedInputLayouts();
+  SmallVector<TensorLayout> initLayouts = opLayout.getSupportedOutputLayouts();
+  std::cout << "Num of input layouts: " << inputLayouts.size() << std::endl;
+  std::cout << "Num of init layouts: " << initLayouts.size() << std::endl;
+
+  // check all inputs and inits are tensor, otherwise no need for layout
+  // propagation
+  bool allTensor =
+      llvm::all_of(inputOperands,
+                   [](OpOperand *opOperand) {
+                     return opOperand->get().getType().isa<TensorType>();
+                   }) &&
+      llvm::all_of(initOperands, [](OpOperand *opOperand) {
+        return opOperand->get().getType().isa<TensorType>();
+      });
+  std::cout << "The op's input is all tensor?" << allTensor << std::endl;
+  if (!allTensor) {
+    return failure("the op does not need packing.");
+  }
+  for (const auto &operandsList : {inputOperands, initOperands}) {
+    for (OpOperand *opOperand : operandsList) {
+      int64_t pos = opOperand->getOperandNumber();
+      std::cout << "pos: " << pos << std::endl;
+      Value operand = opOperand->get();
+      TensorLayout targetLayout = pos >= inputLayouts.size()
+                                      ? initLayouts[pos - inputLayouts.size()]
+                                      : inputLayouts[pos];
+      SmallVector<int64_t> outerPerm = targetLayout.getOuterAxis();
+      SmallVector<int64_t> innerPos = targetLayout.getInnerAxis();
+      SmallVector<OpFoldResult> innerPackSizes = targetLayout.getTileSizes();
+
+      std::cout << "Suggested layout: " << targetLayout << std::endl;
+
+      std::cout << "Operand shape: ";
+      for (auto dim :
+           llvm::cast<RankedTensorType>(operand.getType()).getShape()) {
+        std::cout << dim << ", ";
+      }
+      std::cout << std::endl;
+
+      Value dest = tensor::PackOp::createDestinationTensor(
+          rewriter, loc, operand, innerPackSizes, innerPos, outerPerm);
+      ShapedType operandType = cast<ShapedType>(operand.getType());
+      bool areConstantTiles =
+          llvm::all_of(innerPackSizes, [](OpFoldResult tile) {
+            return getConstantIntValue(tile).has_value();
+          });
+      if (areConstantTiles && operandType.hasStaticShape() &&
+          !tensor::PackOp::requirePaddingValue(
+              operandType.getShape(), innerPos,
+              cast<ShapedType>(dest.getType()).getShape(), {},
+              innerPackSizes)) {
+        packOps.push_back(rewriter.create<tensor::PackOp>(
+            loc, operand, dest, innerPos, innerPackSizes, std::nullopt,
+            outerPerm));
+      } else {
+        // TODO: value of the padding attribute should be determined by
+        // consumers.
+        auto zeroAttr =
+            rewriter.getZeroAttr(getElementTypeOrSelf(dest.getType()));
+        Value zero = rewriter.create<arith::ConstantOp>(loc, zeroAttr);
+        packOps.push_back(rewriter.create<tensor::PackOp>(
+            loc, operand, dest, innerPos, innerPackSizes, zero, outerPerm));
+      }
+      inputsAndInits.push_back(packOps.back());
+    }
+  }
+
+  // Step 3. Build the packed op, use the type of `inits` as result types.
+  ValueRange inputs =
+      ValueRange{inputsAndInits}.take_front(linalgOp.getNumDpsInputs());
+  ValueRange inits =
+      ValueRange{inputsAndInits}.take_back(linalgOp.getNumDpsInits());
+  // TODO(yifei): the axis info of reduce/broadcast/transpose may change
+  auto packedLinalgOp = mlir::clone(
+      rewriter, linalgOp, SmallVector<Type>{inputsAndInits.back().getType()},
+      inputsAndInits);
+
+  // Step 4. Unpack all the op results.
+  for (OpResult result : packedLinalgOp->getResults()) {
+    int64_t resultNum = result.getResultNumber();
+    tensor::PackOp maybePackedInit =
+        inits[resultNum].getDefiningOp<tensor::PackOp>();
+    if (!maybePackedInit) {
+      results.push_back(result);
+      continue;
+    }
+    // Build the symmetrical UnPackOp to the existing PackOp.
+    unPackOps.push_back(rewriter.create<tensor::UnPackOp>(
+        packedLinalgOp->getLoc(), result, maybePackedInit.getSource(),
+        maybePackedInit.getInnerDimsPos(), maybePackedInit.getMixedTiles()));
+    results.push_back(unPackOps.back());
+  }
+
+  // Step 5. Replace `linalgOp`.
+  rewriter.replaceOp(linalgOp, results);
+
+  // Return packedLinalgOp.
+  return linalg::PackResult{
+      packOps, cast<linalg::LinalgOp>(packedLinalgOp.getOperation()),
+      unPackOps};
+}
+
 class PropagateLayout : public impl::PropagateLayoutBase<PropagateLayout> {
 public:
   using impl::PropagateLayoutBase<PropagateLayout>::PropagateLayoutBase;
@@ -42,24 +167,37 @@ void PropagateLayout::runOnOperation() {
   IRRewriter rewriter(ctx);
   // walk the entire graph
   auto &layoutAnalysisResult = getAnalysis<GlobalAnalysis>();
-  graph->walk([&](linalg::LinalgOp linalgOp) {
+  SmallVector<Operation *> packTODOList;
+  graph->walk([&](Operation *op) {
+    if (isa<linalg::LinalgOp>(op) && !mlir::linalg::isaContractionOpInterface(
+                                         dyn_cast<linalg::LinalgOp>(op))) {
+      packTODOList.push_back(op);
+    }
+  });
+  for (auto op : packTODOList) {
     std::cout << std::endl;
     std::cout << "----------------------------------" << std::endl;
     std::cout << "Visiting op ";
-    linalgOp.getOperation()->getName().print(llvm::errs());
+    op->getName().print(llvm::errs());
     std::cout << std::endl;
     std::cout << "----------------------------------" << std::endl;
-    FailureOr<OperatorLayout> opLayout =
-        layoutAnalysisResult.getOpLayout(linalgOp);
+    FailureOr<OperatorLayout> opLayout = layoutAnalysisResult.getOpLayout(op);
     if (failed(opLayout)) {
       std::cout << "infer failed" << std::endl;
     } else {
       // pack op into ideal layout
       std::cout << "-------- supported layouts -------" << std::endl;
       std::cout << *opLayout << std::endl;
       // insert pack
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(op);
+      if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
+        FailureOr<linalg::PackResult> packedOp =
+            packNamedOp(rewriter, linalgOp, *opLayout);
+      }
+      graph->dump();
     }
-  });
+  }
   graph->dump();
 }
 
