[PlutoTransform] fix symbol alignment

Author: kumasento

tools/polymer/lib/Support/ScopStmt.cc

@@ -109,6 +109,17 @@ promoteSymbolToTopLevel(mlir::Value val, FlatAffineValueConstraints &domain,
   symMap[val] = arg;
 }
 
+static void reorderSymbolsByOperandId(FlatAffineValueConstraints &cst) {
+  // bubble sort
+  for (unsigned i = cst.getNumDimIds(); i < cst.getNumDimAndSymbolIds(); ++i)
+    for (unsigned j = i + 1; j < cst.getNumDimAndSymbolIds(); ++j) {
+      auto fst = cst.getValue(i).cast<BlockArgument>();
+      auto snd = cst.getValue(j).cast<BlockArgument>();
+      if (fst.getArgNumber() > snd.getArgNumber())
+        cst.swapId(i, j);
+    }
+}
+
 void ScopStmtImpl::initializeDomainAndEnclosingOps() {
   // Extract the affine for/if ops enclosing the caller and insert them into the
   // enclosingOps list.
@@ -140,6 +151,9 @@ void ScopStmtImpl::initializeDomainAndEnclosingOps() {
                    &symValues);
   for (mlir::Value val : symValues)
     promoteSymbolToTopLevel(val, domain, symMap);
+
+  // Without this things like swapped-bounds.mlir in test cannot work.
+  reorderSymbolsByOperandId(domain);
 }
 
 void ScopStmtImpl::getArgsValueMapping(BlockAndValueMapping &argMap) {

tools/polymer/lib/Target/OpenScop/ConvertToOpenScop.cc

@@ -196,7 +196,33 @@ void OslScopBuilder::buildScopStmtMap(mlir::FuncOp f,
 void OslScopBuilder::buildScopContext(OslScop *scop,
                                       OslScop::ScopStmtMap *scopStmtMap,
                                       FlatAffineValueConstraints &ctx) const {
-  ctx.reset();
+  LLVM_DEBUG(dbgs() << "--- Building SCoP context ...\n");
+
+  // First initialize the symbols of the ctx by the order of arg number.
+  // This simply aims to make mergeAndAlignIdsWithOthers work.
+  SmallVector<Value> symbols;
+  for (const auto &it : *scopStmtMap) {
+    auto domain = it.second.getDomain();
+    SmallVector<Value> syms;
+    domain->getValues(domain->getNumDimIds(), domain->getNumDimAndSymbolIds(),
+                      &syms);
+
+    for (Value sym : syms) {
+      // Find the insertion position.
+      auto it = symbols.begin();
+      while (it != symbols.end()) {
+        auto lhs = it->cast<BlockArgument>();
+        auto rhs = sym.cast<BlockArgument>();
+        if (lhs.getArgNumber() >= rhs.getArgNumber())
+          break;
+        ++it;
+      }
+      if (*it != sym)
+        symbols.insert(it, sym);
+    }
+  }
+  ctx.reset(/*numDims=*/0, /*numSymbols=*/symbols.size());
+  ctx.setValues(0, symbols.size(), symbols);
 
   // Union with the domains of all Scop statements. We first merge and align the
   // IDs of the context and the domain of the scop statement, and then append
@@ -210,6 +236,31 @@ void OslScopBuilder::buildScopContext(OslScop *scop,
     ctx.mergeAndAlignIdsWithOther(0, &cst);
     ctx.append(cst);
     ctx.removeRedundantConstraints();
+
+    LLVM_DEBUG(dbgs() << "Statement:\n");
+    LLVM_DEBUG(it.second.getCaller().dump());
+    LLVM_DEBUG(it.second.getCallee().dump());
+    LLVM_DEBUG(dbgs() << "Target domain: \n");
+    LLVM_DEBUG(domain->dump());
+
+    LLVM_DEBUG({
+      dbgs() << "Domain values: \n";
+      SmallVector<Value> values;
+      domain->getValues(0, domain->getNumDimAndSymbolIds(), &values);
+      for (Value value : values)
+        dbgs() << " * " << value << '\n';
+    });
+
+    LLVM_DEBUG(dbgs() << "Updated context: \n");
+    LLVM_DEBUG(ctx.dump());
+
+    LLVM_DEBUG({
+      dbgs() << "Context values: \n";
+      SmallVector<Value> values;
+      ctx.getValues(0, ctx.getNumDimAndSymbolIds(), &values);
+      for (Value value : values)
+        dbgs() << " * " << value << '\n';
+    });
   }
 
   // Then, create the single context relation in scop.
@@ -221,19 +272,43 @@ void OslScopBuilder::buildScopContext(OslScop *scop,
   SmallVector<mlir::Value, 8> symValues;
   ctx.getValues(ctx.getNumDimIds(), ctx.getNumDimAndSymbolIds(), &symValues);
 
+  // Add and align domain SYMBOL columns.
   for (const auto &it : *scopStmtMap) {
     FlatAffineValueConstraints *domain = it.second.getDomain();
+    // For any symbol missing in the domain, add them directly to the end.
+    for (unsigned i = 0; i < ctx.getNumSymbolIds(); ++i) {
+      unsigned pos;
+      if (!domain->findId(symValues[i], &pos)) // insert to the back
+        domain->appendSymbolId(symValues[i]);
+      else
+        LLVM_DEBUG(dbgs() << "Found " << symValues[i] << '\n');
+    }
 
+    // Then do the aligning.
+    LLVM_DEBUG(domain->dump());
     for (unsigned i = 0; i < ctx.getNumSymbolIds(); i++) {
       mlir::Value sym = symValues[i];
       unsigned pos;
-      if (domain->findId(sym, &pos)) {
-        if (pos != i + domain->getNumDimIds())
-          domain->swapId(i + domain->getNumDimIds(), pos);
-      } else {
-        domain->insertSymbolId(i, sym);
-      }
+      assert(domain->findId(sym, &pos));
+
+      unsigned posAsCtx = i + domain->getNumDimIds();
+      LLVM_DEBUG(dbgs() << "Swapping " << posAsCtx << " " << pos << "\n");
+      if (pos != posAsCtx)
+        domain->swapId(posAsCtx, pos);
     }
+
+    // for (unsigned i = 0; i < ctx.getNumSymbolIds(); i++) {
+    //   mlir::Value sym = symValues[i];
+    //   unsigned pos;
+    //   // If the symbol can be found in the domain, we put it in the same
+    //   // position as the ctx.
+    //   if (domain->findId(sym, &pos)) {
+    //     if (pos != i + domain->getNumDimIds())
+    //       domain->swapId(i + domain->getNumDimIds(), pos);
+    //   } else {
+    //     domain->insertSymbolId(i, sym);
+    //   }
+    // }
   }
 }
 

tools/polymer/lib/Transforms/PlutoTransform.cc

@@ -37,6 +37,8 @@ using namespace mlir;
 using namespace llvm;
 using namespace polymer;
 
+#define DEBUG_TYPE "pluto-opt"
+
 namespace {
 struct PlutoOptPipelineOptions
     : public mlir::PassPipelineOptions<PlutoOptPipelineOptions> {
@@ -72,6 +74,8 @@ static mlir::FuncOp plutoTransform(mlir::FuncOp f, OpBuilder &rewriter,
                                    bool parallelize = false, bool debug = false,
                                    int cloogf = -1, int cloogl = -1,
                                    bool diamondTiling = false) {
+  LLVM_DEBUG(dbgs() << "Pluto transforming: \n");
+  LLVM_DEBUG(f.dump());
 
   PlutoContext *context = pluto_context_alloc();
   OslSymbolTable srcTable, dstTable;
@@ -128,6 +132,28 @@ static mlir::FuncOp plutoTransform(mlir::FuncOp f, OpBuilder &rewriter,
   return g;
 }
 
+static void dedupIndexCast(FuncOp f) {
+  Block &entry = f.getBlocks().front();
+  llvm::MapVector<Value, Value> argToCast;
+  SmallVector<Operation *> toErase;
+  for (auto &op : entry) {
+    if (auto indexCast = dyn_cast<arith::IndexCastOp>(&op)) {
+      auto arg = indexCast.getOperand().dyn_cast<BlockArgument>();
+      if (argToCast.count(arg)) {
+        LLVM_DEBUG(dbgs() << "Found duplicated index_cast: " << indexCast
+                          << '\n');
+        indexCast.replaceAllUsesWith(argToCast.lookup(arg));
+        toErase.push_back(indexCast);
+      } else {
+        argToCast[arg] = indexCast;
+      }
+    }
+  }
+
+  for (auto op : toErase)
+    op->erase();
+}
+
 namespace {
 class PlutoTransformPass
     : public mlir::PassWrapper<PlutoTransformPass,
@@ -156,8 +182,10 @@ class PlutoTransformPass
     llvm::DenseMap<mlir::FuncOp, mlir::FuncOp> funcMap;
 
     m.walk([&](mlir::FuncOp f) {
-      if (!f->getAttr("scop.stmt") && !f->hasAttr("scop.ignored"))
+      if (!f->getAttr("scop.stmt") && !f->hasAttr("scop.ignored")) {
+        dedupIndexCast(f);
         funcOps.push_back(f);
+      }
     });
 
     for (mlir::FuncOp f : funcOps)

tools/polymer/test/polymer-opt/PlutoTransforms/swapped-bounds.mlir

@@ -0,0 +1,21 @@
+// RUN: polymer-opt %s -pluto-opt | FileCheck %s
+func private @S0() attributes {scop.stmt}
+func private @S1() attributes {scop.stmt}
+
+func @foo(%N: index, %M: index, %L: index) {
+  affine.for %i = 0 to %N {
+    affine.for %j = 0 to %L {
+      call @S0() : () -> ()
+    }
+    affine.for %j = 0 to %M {
+      affine.for %k = 0 to %L {
+        call @S1() : () -> ()
+      }
+    }
+  }
+  return
+}
+
+
+// Just need to check this thing can be transformed.
+// CHECK: func @foo