Utility function to map a loop on a parametric grid of virtual processors

This CL introduces a simple loop utility function which rewrites the bounds and step of a loop so as to become mappable on a regular grid of processors whose identifiers are given by SSA values.

A corresponding unit test is added.

For example, using CUDA terminology, and assuming a 2-d grid with processorIds = [blockIdx.x, threadIdx.x] and numProcessors = [gridDim.x, blockDim.x], the loop:
```
   loop.for %i = %lb to %ub step %step {
     ...
   }
```
is rewritten into a version resembling the following pseudo-IR:
```
   loop.for %i = %lb + threadIdx.x + blockIdx.x * blockDim.x to %ub
      step %gridDim.x * blockDim.x {
     ...
   }
```

PiperOrigin-RevId: 258945942

1 files changed, 65 insertions, 0 deletions

diff --git a/mlir/test/lib/Transforms/TestLoopMapping.cpp b/mlir/test/lib/Transforms/TestLoopMapping.cpp
new file mode 100644
index 00000000000..fb1ef64d26f
--- /dev/null
+++ b/mlir/test/lib/Transforms/TestLoopMapping.cpp
@@ -0,0 +1,65 @@
+//===- TestLoopMapping.cpp --- Parametric loop mapping pass ---------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements a pass to parametrically map loop.for loops to virtual
+// processing element dimensions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LoopOps/LoopOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+
+namespace {
+class TestLoopMappingPass : public FunctionPass<TestLoopMappingPass> {
+public:
+  explicit TestLoopMappingPass() {}
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+
+    // SSA values for the transformation are created out of thin air by
+    // unregistered "new_processor_id_and_range" operations. This is enough to
+    // emulate mapping conditions.
+    SmallVector<Value *, 8> processorIds, numProcessors;
+    func.walk([&processorIds, &numProcessors](Operation *op) {
+      if (op->getName().getStringRef() != "new_processor_id_and_range")
+        return;
+      processorIds.push_back(op->getResult(0));
+      numProcessors.push_back(op->getResult(1));
+    });
+
+    func.walk<loop::ForOp>([&processorIds, &numProcessors](loop::ForOp op) {
+      // Ignore nested loops.
+      if (op.getContainingRegion()->getParentOfType<loop::ForOp>())
+        return;
+      mapLoopToProcessorIds(op, processorIds, numProcessors);
+    });
+  }
+};
+} // end namespace
+
+static PassRegistration<TestLoopMappingPass>
+    reg("test-mapping-to-processing-elements",
+        "test mapping a single loop on a virtual processor grid",
+        [] { return new TestLoopMappingPass(); });