[PlutoTransform] can handle index at top-level

Author: kumasento

tools/polymer/lib/Support/OslScop.cc

@@ -585,6 +585,12 @@ void OslScop::createAccessRelationConstraints(
   cst.reset();
   cst.mergeAndAlignIdsWithOther(0, &domain);
 
+  LLVM_DEBUG({
+    dbgs() << "Building access relation.\n"
+           << " + Domain:\n";
+    domain.dump();
+  });
+
   SmallVector<mlir::Value, 8> idValues;
   domain.getAllValues(&idValues);
   llvm::SetVector<mlir::Value> idValueSet;

tools/polymer/lib/Support/ScopStmt.cc

@@ -115,7 +115,20 @@ void ScopStmtImpl::initializeDomainAndEnclosingOps() {
   getEnclosingAffineForAndIfOps(*caller, &enclosingOps);
 
   // The domain constraints can then be collected from the enclosing ops.
-  getIndexSet(enclosingOps, &domain);
+  assert(succeeded(getIndexSet(enclosingOps, &domain)));
+
+  // Add additional indices that are in the top level block arguments.
+  for (Value arg : caller->getOperands()) {
+    if (!arg.getType().isIndex())
+      continue;
+    unsigned pos;
+    if (domain.findId(arg, &pos))
+      continue;
+
+    domain.appendSymbolId(1);
+    domain.dump();
+    domain.setValue(domain.getNumDimAndSymbolIds() - 1, arg);
+  }
 
   // Symbol values, which could be a BlockArgument, or the result of DimOp or
   // IndexCastOp, or even an affine.apply. Here we limit the cases to be either
@@ -181,6 +194,8 @@ static mlir::Value findBlockArg(mlir::Value v) {
 void ScopStmt::getAccessMapAndMemRef(mlir::Operation *op,
                                      mlir::AffineValueMap *vMap,
                                      mlir::Value *memref) const {
+  // Map from callee arguments to caller's. impl holds the callee and caller
+  // instances.
   BlockAndValueMapping argMap;
   impl->getArgsValueMapping(argMap);
 

tools/polymer/test/polymer-opt/PlutoTransforms/subindex.mlir

@@ -0,0 +1,20 @@
+// RUN: polymer-opt %s -pluto-opt | FileCheck %s
+
+
+func @S0(%A: memref<10x60xf32>, %i: index, %j: index) attributes {scop.stmt} {
+  %cst = arith.constant 3.14 : f32
+  affine.store %cst, %A[%i, %j] : memref<10x60xf32>
+  return
+}
+
+func @foo(%A: memref<10x60xf32>, %i: index) {
+  affine.for %j = 0 to 60 {
+    call @S0(%A, %i, %j): (memref<10x60xf32>, index, index) -> ()
+  }
+  return
+}
+
+// CHECK: func @foo(%[[A:.*]]: memref<{{.*}}>, %[[i:.*]]: index)
+// CHECK-NEXT: affine.for 
+// CHECK-NEXT: affine.for %[[j:.*]] =
+// CHECK-NEXT: call @S0(%[[A]], %[[i]], %[[j]])