44 lines · cpp
1//===- GlobalIdRewriter.cpp - Implementation of GlobalId rewriting -------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements in-dialect rewriting of the global_id op for archs10// where global_id.x = threadId.x + blockId.x * blockDim.x11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/GPU/IR/GPUDialect.h"15#include "mlir/Dialect/GPU/Transforms/Passes.h"16#include "mlir/Dialect/Index/IR/IndexOps.h"17#include "mlir/IR/PatternMatch.h"18 19using namespace mlir;20 21namespace {22struct GpuGlobalIdRewriter : public OpRewritePattern<gpu::GlobalIdOp> {23 using OpRewritePattern<gpu::GlobalIdOp>::OpRewritePattern;24 25 LogicalResult matchAndRewrite(gpu::GlobalIdOp op,26 PatternRewriter &rewriter) const override {27 Location loc = op.getLoc();28 auto dim = op.getDimension();29 auto blockId = gpu::BlockIdOp::create(rewriter, loc, dim);30 auto blockDim = gpu::BlockDimOp::create(rewriter, loc, dim);31 // Compute blockId.x * blockDim.x32 auto tmp = index::MulOp::create(rewriter, op.getLoc(), blockId, blockDim);33 auto threadId = gpu::ThreadIdOp::create(rewriter, loc, dim);34 // Compute threadId.x + blockId.x * blockDim.x35 rewriter.replaceOpWithNewOp<index::AddOp>(op, threadId, tmp);36 return success();37 }38};39} // namespace40 41void mlir::populateGpuGlobalIdPatterns(RewritePatternSet &patterns) {42 patterns.add<GpuGlobalIdRewriter>(patterns.getContext());43}44