add logic for reduce and broadcast

Author: yifeizh2

lib/gc/Analysis/GlobalAnalysis.cpp

@@ -1,5 +1,4 @@
-//===- GlobalAnalysis.cpp - Propagate pack unpack on linalg named ops --*- C++
-//-*-===//
+//===- GlobalAnalysis.cpp - Propagate packing on linalg named ops *- C++-*-===//
 //
 // This file is only temporarily used to extend upstream or upcoming utility in
 // TilingInterface, which finally aims for upstream.
@@ -101,7 +100,8 @@ inferIndexingMapRelation(AffineMap indexingMapBase,
 
 // given j --> i and max rank of i, return i --> j
 static DenseMap<int64_t, int64_t>
-getReversedIndexMap(DenseMap<int64_t, int64_t> indexMap, size_t maxRank) {
+getReversedIndexMap(const DenseMap<int64_t, int64_t> &indexMap,
+                    size_t maxRank) {
   DenseMap<int64_t, int64_t> res;
   for (auto pair : indexMap) {
     if (pair.second >= 0) {
@@ -118,7 +118,7 @@ getReversedIndexMap(DenseMap<int64_t, int64_t> indexMap, size_t maxRank) {
 
 static FailureOr<TensorLayout>
 inferTargetLayout(TensorLayout layoutBase,
-                  DenseMap<int64_t, int64_t> indexMap) {
+                  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();

lib/gc/Transforms/PropagateLayout.cpp

@@ -31,6 +31,19 @@ using namespace mlir;
 using namespace mlir::arith;
 using namespace mlir::tensor;
 
+static SmallVector<int64_t> getPackedAxes(ArrayRef<int64_t> dimensions,
+                                          TensorLayout targetLayout) {
+  SmallVector<int64_t> result(dimensions);
+  auto innerPos = targetLayout.getInnerAxis();
+  for (size_t i = 0; i < dimensions.size(); ++i) {
+    if (std::find(innerPos.begin(), innerPos.end(), dimensions[i]) !=
+        innerPos.end()) {
+      result.push_back(i + targetLayout.getOuterAxis().size());
+    }
+  }
+  return result;
+}
+
 static FailureOr<linalg::PackResult> packNamedOp(RewriterBase &rewriter,
                                                  linalg::LinalgOp linalgOp,
                                                  OperatorLayout opLayout) {
@@ -125,10 +138,30 @@ static FailureOr<linalg::PackResult> packNamedOp(RewriterBase &rewriter,
       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);
+  // TODO(yifei): deal with reduce/broadcast/transpose
+  // TODO(yifei): deal with generic
+  linalg::LinalgOp packedLinalgOp;
+  if (auto reduceOp = dyn_cast<linalg::ReduceOp>(&linalgOp)) {
+    SmallVector<int64_t> packedAxes =
+        getPackedAxes(reduceOp->getDimensions(), inputLayouts[0]);
+    packedLinalgOp = rewriter.create<linalg::ReduceOp>(
+        loc, inits.getTypes(), inputs, inits, packedAxes);
+    packedLinalgOp->getRegion(0).takeBody(linalgOp->getRegion(0));
+  } else if (auto broadcastOp = dyn_cast<linalg::BroadcastOp>(&linalgOp)) {
+    packedLinalgOp = rewriter.create<linalg::BroadcastOp>(
+        loc, inputs[0], inits[0], broadcastOp->getDimensions());
+  } else if (isa<linalg::TransposeOp>(linalgOp)) {
+    // remove transpose op
+  } else if (isa<linalg::SoftmaxOp>(linalgOp) ||
+             isa<linalg::GenericOp>(linalgOp) || isa<linalg::MapOp>(linalgOp) ||
+             isa<linalg::YieldOp>(linalgOp) || isa<linalg::IndexOp>(linalgOp)) {
+    return failure(
+        "Packing logic not implemented for SoftMax/Generic/Map/Yield/Index.");
+  } else {
+    packedLinalgOp = mlir::clone(
+        rewriter, linalgOp, SmallVector<Type>{inputsAndInits.back().getType()},
+        inputsAndInits);
+  }
 
   // Step 4. Unpack all the op results.
   for (OpResult result : packedLinalgOp->getResults()) {
@@ -195,7 +228,6 @@ void PropagateLayout::runOnOperation() {
         FailureOr<linalg::PackResult> packedOp =
             packNamedOp(rewriter, linalgOp, *opLayout);
       }
-      graph->dump();
     }
   }
   graph->dump();