support expand shape

yifeizh2 · yifeizh2 · commit ca0e2b6c84a3 · 2024-06-02T23:33:04.000-07:00
diff --git a/include/gc/Analysis/GlobalAnalysis.h b/include/gc/Analysis/GlobalAnalysis.h
@@ -18,6 +18,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
+#include <llvm/Support/Debug.h>
 
 namespace mlir {
 namespace gc {
@@ -27,17 +28,9 @@ using namespace mlir;
 class TensorLayout {
 public:
   TensorLayout(ArrayRef<int64_t> outerAxis, ArrayRef<int64_t> innerAxis,
-               ArrayRef<OpFoldResult> tileSizes) {
+               ArrayRef<OpFoldResult> tileSizes)
+      : OuterAxis(outerAxis), InnerAxis(innerAxis), TileSizes(tileSizes) {
     assert(innerAxis.size() == tileSizes.size());
-    for (auto oa : outerAxis) {
-      OuterAxis.push_back(oa);
-    }
-    for (auto ia : innerAxis) {
-      InnerAxis.push_back(ia);
-    }
-    for (auto ts : tileSizes) {
-      TileSizes.push_back(ts);
-    }
   }
 
   bool isPlainLayout() const {
@@ -55,25 +48,22 @@ class TensorLayout {
                         SmallVector<OpFoldResult>{});
   }
 
-  static DenseMap<int64_t, SmallVector<int64_t>>
-  getPlain2PackedMapping(TensorLayout layout) {
+  DenseMap<int64_t, SmallVector<int64_t>> getPlain2PackedMapping() {
     DenseMap<int64_t, SmallVector<int64_t>> p2b;
-    SmallVector<int64_t> outerAxis = layout.getOuterAxis();
-    SmallVector<int64_t> innerAxis = layout.getInnerAxis();
-    for (size_t i = 0; i < outerAxis.size(); ++i) {
-      p2b[outerAxis[i]].push_back(i);
+    for (size_t i = 0; i < OuterAxis.size(); ++i) {
+      p2b[OuterAxis[i]].push_back(i);
     }
-    for (size_t i = 0; i < innerAxis.size(); ++i) {
-      p2b[innerAxis[i]].push_back(outerAxis.size() + i);
+    for (size_t i = 0; i < InnerAxis.size(); ++i) {
+      p2b[InnerAxis[i]].push_back(InnerAxis.size() + i);
     }
     return p2b;
   }
 
   FailureOr<int64_t> getOriginalAxis(int64_t idx) {
-    size_t totalRank = OuterAxis.size() + InnerAxis.size();
+    int64_t totalRank = OuterAxis.size() + InnerAxis.size();
     if (idx >= totalRank) {
       return failure("Index out of range.");
-    } else if (idx >= OuterAxis.size()) {
+    } else if (idx >= static_cast<int64_t>(OuterAxis.size())) {
       return InnerAxis[idx - OuterAxis.size()];
     } else {
       return OuterAxis[idx];
@@ -88,7 +78,8 @@ class TensorLayout {
 
   SmallVector<OpFoldResult> getTileSizes() const { return TileSizes; }
 
-  friend std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout);
+  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
+                                       const TensorLayout &layout);
 
   bool operator==(const TensorLayout &layout);
 
@@ -121,8 +112,8 @@ class OperatorLayout {
     return supportedOutputLayouts[idx];
   }
 
-  friend std::ostream &operator<<(std::ostream &ss,
-                                  const OperatorLayout &opLayout);
+  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
+                                       const OperatorLayout &opLayout);
 
 private:
   SmallVector<TensorLayout> supportedInputLayouts;
@@ -134,14 +125,14 @@ class GlobalAnalysis {
   explicit GlobalAnalysis(Operation *root);
 
   FailureOr<OperatorLayout> getOpLayout(Operation *op) {
-    if (layout.find(op) != layout.end())
-      return layout[op];
+    if (layoutCache.find(op) != layoutCache.end())
+      return layoutCache[op];
     else
       return failure("Current op does not have layout information.");
   }
 
 private:
-  DenseMap<Operation *, OperatorLayout> layout;
+  DenseMap<Operation *, OperatorLayout> layoutCache;
 };
 
 } // namespace gc
diff --git a/include/gc/Transforms/Passes.td b/include/gc/Transforms/Passes.td
@@ -44,12 +44,13 @@ def GCCPUPipeline: Pass<"gc-cpu-pipeline"> {
       "vector::VectorDialect"];
 }
 
-def PropagateLayout : Pass<"propagate-layout"> {
-  let summary = "Insert and propagte tensor.pack to pack the computation of general linalg named ops and tensor ops.";
+def PropagateLayoutOnNamedOps : Pass<"propagate-layout-on-named-ops"> {
+  let summary = "Insert and propagte tensor.pack to pack the computation of linalg named ops and tensor ops.";
   let description = [{
     Insert and propagte tensor.pack
   }];
-  let dependentDialects = ["mlir::tensor::TensorDialect", "mlir::linalg::LinalgDialect"];
+  let dependentDialects = ["mlir::tensor::TensorDialect",
+      "mlir::linalg::LinalgDialect"];
 }
 
 #endif // GC_DIALECT_GC_PASSES
diff --git a/include/gc/Transforms/Transforms.h b/include/gc/Transforms/Transforms.h
@@ -0,0 +1,28 @@
+//===- Transforms.h - transformation utilities ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GC_TRANSFORMS_TRANSFORMS_H
+#define GC_TRANSFORMS_TRANSFORMS_H
+
+#include "gc/Analysis/GlobalAnalysis.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+
+namespace mlir {
+namespace gc {
+FailureOr<linalg::PackResult> packNamedOp(RewriterBase &rewriter,
+                                          linalg::LinalgOp linalgOp,
+                                          OperatorLayout opLayout);
+
+LogicalResult namedOpLayoutPropagation(RewriterBase &rewriter,
+                                       linalg::LinalgOp linalgOp,
+                                       OperatorLayout opLayout);
+} // namespace gc
+} // namespace mlir
+
+#endif // GC_TRANSFORMS_TRANSFORMS_H
diff --git a/lib/gc/Analysis/GlobalAnalysis.cpp b/lib/gc/Analysis/GlobalAnalysis.cpp
@@ -12,10 +12,13 @@
 namespace mlir {
 namespace gc {
 
-std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout) {
-  SmallVector<int64_t> outerAxis = layout.getOuterAxis();
-  SmallVector<int64_t> innerAxis = layout.getInnerAxis();
-  SmallVector<OpFoldResult> tileSizes = layout.getTileSizes();
+#define DEBUG_TYPE "global-analysis"
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
+                              const TensorLayout &layoutCache) {
+  SmallVector<int64_t> outerAxis = layoutCache.getOuterAxis();
+  SmallVector<int64_t> innerAxis = layoutCache.getInnerAxis();
+  SmallVector<OpFoldResult> tileSizes = layoutCache.getTileSizes();
   ss << "[";
   for (size_t i = 0; i < outerAxis.size(); ++i) {
     if (i != 0) {
@@ -43,21 +46,21 @@ std::ostream &operator<<(std::ostream &ss, const TensorLayout &layout) {
   return ss;
 }
 
-bool TensorLayout::operator==(const TensorLayout &layout) {
-  return (this->OuterAxis == layout.getOuterAxis()) &&
-         (this->InnerAxis == layout.getInnerAxis()) &&
-         (this->TileSizes == layout.getTileSizes());
+bool TensorLayout::operator==(const TensorLayout &layoutCache) {
+  return (this->OuterAxis == layoutCache.getOuterAxis()) &&
+         (this->InnerAxis == layoutCache.getInnerAxis()) &&
+         (this->TileSizes == layoutCache.getTileSizes());
 }
 
-std::ostream &operator<<(std::ostream &ss, const OperatorLayout &opLayout) {
-  ss << "operator has " << opLayout.getSupportedInputLayouts().size()
-     << " inputs; " << opLayout.getSupportedOutputLayouts().size()
-     << " outputs." << std::endl;
-  for (const auto &layout : opLayout.getSupportedInputLayouts()) {
-    ss << "input layout: " << layout << std::endl;
+llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
+                              const OperatorLayout &opLayout) {
+  for (auto &&[idx, layoutCache] :
+       llvm::enumerate(opLayout.getSupportedInputLayouts())) {
+    ss << "input " << idx << "'s layoutCache: " << layoutCache << "\n";
   }
-  for (const auto &layout : opLayout.getSupportedOutputLayouts()) {
-    ss << "output layout: " << layout << std::endl;
+  for (auto &&[idx, layoutCache] :
+       llvm::enumerate(opLayout.getSupportedOutputLayouts())) {
+    ss << "output " << idx << "'s layoutCache: " << layoutCache << "\n";
   }
   return ss;
 }
@@ -119,7 +122,6 @@ getReversedIndexMap(const DenseMap<int64_t, int64_t> &indexMap,
 static FailureOr<TensorLayout>
 inferTargetLayout(TensorLayout layoutBase,
                   const DenseMap<int64_t, int64_t> &indexMap) {
-  int64_t dimDifference = indexMap.size() - layoutBase.getTensorRank();
   SmallVector<int64_t> baseOuterAxis = layoutBase.getOuterAxis();
   SmallVector<int64_t> baseInnerAxis = layoutBase.getInnerAxis();
   SmallVector<OpFoldResult> baseTileSizes = layoutBase.getTileSizes();
@@ -153,38 +155,24 @@ inferTargetLayout(TensorLayout layoutBase,
 
 GlobalAnalysis::GlobalAnalysis(Operation *root) {
   root->walk([&](Operation *op) {
+    // get input layouts
+    LLVM_DEBUG(llvm::dbgs()
+               << "Inferring layoutCache of op: " << op->getName() << "\n");
     if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
-      // get input layouts
-      std::cout << std::endl;
-      std::cout << "----------------------------------" << std::endl;
-      linalgOp.getOperation()->getName().print(llvm::errs());
-      std::cout << std::endl;
-      std::cout << "----------------------------------" << std::endl;
-      std::cout << std::endl;
-      SmallVector<AffineMap> indexing_maps = linalgOp.getIndexingMapsArray();
       auto curInputs = linalgOp.getDpsInputOperands();
       auto curResults = linalgOp.getOperation()->getResults();
-
       // ---------------- Get Current Input Layouts -------------------
-      // get current input layouts
-      std::cout << "----- printing ground-truth input layouts -----"
-                << std::endl;
       SmallVector<TensorLayout> curInputLayouts;
       for (auto input : curInputs) {
         auto parent = input->get().getDefiningOp();
-        if (layout.find(parent) != layout.end()) {
+        if (layoutCache.find(parent) != layoutCache.end()) {
           // TODO(yifei): it is not always 0 here
-          curInputLayouts.push_back(layout[parent].getOutputLayout(0));
+          curInputLayouts.push_back(layoutCache[parent].getOutputLayout(0));
         } else {
           curInputLayouts.push_back(TensorLayout::createPlainLayout(
               linalgOp.getMatchingIndexingMap(input).getNumResults()));
         }
       }
-      // debug info
-      for (auto layout : curInputLayouts) {
-        std::cout << "layout: " << layout << std::endl;
-      }
-
       // ------ Get Current Op's Suggested Layout & Do Propagation ------
       IRRewriter rewriter(linalgOp);
       if (mlir::linalg::isaContractionOpInterface(linalgOp)) {
@@ -193,38 +181,33 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
         // curInputLayouts);
 
         // hardcode one for now
-        // A side layout, [0, 1, 0, 1]; {32, 32}
+        // A side layoutCache, [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}
+        // B side layoutCache, [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}
+        // C side layoutCache, [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;
+        layoutCache[linalgOp] = suggestedLayout;
       } else {
         SmallVector<TensorLayout> inputLayouts, outputLayouts;
         inputLayouts.push_back(curInputLayouts[0]);
         // TODO(yifei): wisely choose the input format basis
         // Let's only refer to input[0] for now
         for (size_t i = 1; i < curInputs.size(); ++i) {
-          std::cout << "inferring indexing map relation" << std::endl;
           // getMatchingIndexingMap
           auto res = inferIndexingMapRelation(
               linalgOp.getMatchingIndexingMap(curInputs[0]),
               linalgOp.getMatchingIndexingMap(curInputs[i]));
-          for (auto tp : *res) {
-            std::cout << "target index: " << tp.first
-                      << " maps to base index: " << tp.second << std::endl;
-          }
           TensorLayout inputLayout =
               *inferTargetLayout(curInputLayouts[0], *res);
           inputLayouts.push_back(inputLayout);
@@ -235,14 +218,66 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
         TensorLayout outputLayout =
             *inferTargetLayout(curInputLayouts[0], *res_out);
         outputLayouts.push_back(outputLayout);
-        for (auto tp : *res_out) {
-          std::cout << "target index: " << tp.first
-                    << " maps to base index: " << tp.second << std::endl;
-        }
         OperatorLayout suggestedLayout(inputLayouts, outputLayouts);
-        layout[linalgOp] = suggestedLayout;
+        layoutCache[linalgOp] = suggestedLayout;
+      }
+    } else if (auto padOp = dyn_cast<tensor::PadOp>(op)) {
+      auto inputOperand = padOp.getSource();
+      auto inputRank =
+          cast<ShapedType>(inputOperand.getType()).getShape().size();
+      auto parent = inputOperand.getDefiningOp();
+      TensorLayout curInputLayout =
+          layoutCache.find(parent) != layoutCache.end()
+              ? layoutCache[parent].getOutputLayout(0)
+              : TensorLayout::createPlainLayout(inputRank);
+      SmallVector<TensorLayout> inputLayouts{curInputLayout},
+          outputLayouts{curInputLayout};
+      OperatorLayout suggestedLayout(inputLayouts, outputLayouts);
+      layoutCache[padOp] = suggestedLayout;
+    } else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(op)) {
+      auto reassociation = expandShapeOp.getReassociation();
+      auto staticOutputShape = expandShapeOp.getStaticOutputShape();
+      auto parent = expandShapeOp.getSrc().getDefiningOp();
+      auto inputShape = expandShapeOp.getSrcType().getShape();
+      TensorLayout curInputLayout =
+          layoutCache.find(parent) != layoutCache.end()
+              ? layoutCache[parent].getOutputLayout(0)
+              : TensorLayout::createPlainLayout(inputShape.size());
+      DenseMap<int64_t, int64_t> outputInputIdxMapping, inputOutputIndexMapping;
+      int64_t accumulationOffset = 0;
+      for (int64_t i = 0; i < static_cast<int64_t>(reassociation.size()); ++i) {
+        auto subReassociation = llvm::cast<ArrayAttr>(reassociation[i]);
+        for (int64_t j = 0; j < static_cast<int64_t>(subReassociation.size());
+             ++j) {
+          if (staticOutputShape[accumulationOffset + j] == inputShape[i]) {
+            outputInputIdxMapping[accumulationOffset + j] = i;
+            inputOutputIndexMapping[i] = accumulationOffset + j;
+          }
+        }
+        accumulationOffset += subReassociation.size();
+      }
+      auto inputOuterAxis = curInputLayout.getOuterAxis();
+      auto inputInnerAxis = curInputLayout.getInnerAxis();
+      int64_t startIdx = 0;
+      SmallVector<int64_t> outputOuterAxis, outputInnerAxis;
+      for (int64_t i = 0; i < static_cast<int64_t>(staticOutputShape.size());
+           ++i) {
+        if (outputInputIdxMapping.find(i) != outputInputIdxMapping.end()) {
+          outputOuterAxis.push_back(inputOuterAxis[outputInputIdxMapping[i]]);
+        } else {
+          outputOuterAxis.push_back(startIdx++);
+        }
+      }
+      for (int64_t i = 0; i < static_cast<int64_t>(inputInnerAxis.size());
+           ++i) {
+        outputInnerAxis.push_back(inputOutputIndexMapping[inputInnerAxis[i]]);
       }
-    } else if (isa<tensor::PadOp>(op) || isa<tensor::ExpandShapeOp>(op)) {
+      TensorLayout outputLayout(outputOuterAxis, outputInnerAxis,
+                                curInputLayout.getTileSizes());
+      SmallVector<TensorLayout> inputLayouts{curInputLayout},
+          outputLayouts{outputLayout};
+      OperatorLayout suggestedLayout(inputLayouts, outputLayouts);
+      layoutCache[expandShapeOp] = suggestedLayout;
     }
   });
 }
diff --git a/lib/gc/Transforms/PropagateLayout.cpp b/lib/gc/Transforms/PropagateLayout.cpp
