db4cd1c8dc
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
66 lines
2.3 KiB
C++
66 lines
2.3 KiB
C++
//===- 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(); });
|