llvm/mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp

//===- DecomposeMemRefs.cpp - Decompose memrefs pass implementation -------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements decompose memrefs pass.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir {
#define GEN_PASS_DEF_GPUDECOMPOSEMEMREFSPASS
#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
} // namespace mlir

usingnamespacemlir;

static void setInsertionPointToStart(OpBuilder &builder, Value val) { … }

static bool isInsideLaunch(Operation *op) { … }

static std::tuple<Value, OpFoldResult, SmallVector<OpFoldResult>>
getFlatOffsetAndStrides(OpBuilder &rewriter, Location loc, Value source,
                        ArrayRef<OpFoldResult> subOffsets,
                        ArrayRef<OpFoldResult> subStrides = std::nullopt) { … }

static Value getFlatMemref(OpBuilder &rewriter, Location loc, Value source,
                           ValueRange offsets) { … }

static bool needFlatten(Value val) { … }

static bool checkLayout(Value val) { … }

namespace {
struct FlattenLoad : public OpRewritePattern<memref::LoadOp> { … };

struct FlattenStore : public OpRewritePattern<memref::StoreOp> { … };

struct FlattenSubview : public OpRewritePattern<memref::SubViewOp> { … };

struct GpuDecomposeMemrefsPass
    : public impl::GpuDecomposeMemrefsPassBase<GpuDecomposeMemrefsPass> { … };

} // namespace

void mlir::populateGpuDecomposeMemrefsPatterns(RewritePatternSet &patterns) { … }

std::unique_ptr<Pass> mlir::createGpuDecomposeMemrefsPass() { … }