llvm/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

//===- TestTensorTransforms.cpp - Test Tensor transformation patterns -----===//
//
// 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 logic for testing Tensor transformations.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Transform/IR/TransformOps.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

usingnamespacemlir;

namespace {
struct TestTensorTransforms
    : public PassWrapper<TestTensorTransforms, OperationPass<>> { … };
} // namespace

static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) { … }

static void applyBubbleUpExpandShapePatterns(Operation *rootOp) { … }

static void applyFoldIntoPackAndUnpackPatterns(Operation *rootOp) { … }

static void applyFoldConstantExtractSlicePatterns(Operation *rootOp) { … }

static void applyFoldConsecutiveInsertExtractSlicePatterns(Operation *rootOp) { … }

static void
applyDropRedundantInsertSliceRankExpansionPatterns(Operation *rootOp) { … }

static void applySimplifyPackUnpackPatterns(Operation *rootOp) { … }

namespace {
/// Base pattern to rewrite  a `tensor.collapse_shape -> tensor.extract_slice`.
/// The `tensor.extract_slice` is replaced by a loop or gather operation that
/// stitches together the desired tile from slices of the source of the collapse
/// shape op.
struct RewriteExtractSliceFromCollapseShapeBase
    : public OpRewritePattern<tensor::ExtractSliceOp> { … };

struct RewriteExtractSliceFromCollapseShapeUsingScfFor
    : public RewriteExtractSliceFromCollapseShapeBase { … };

struct RewriteExtractSliceFromCollapseShapeUsingScfForeach
    : public RewriteExtractSliceFromCollapseShapeBase { … };
} // namespace

static LogicalResult
applyRewriteExtractFromCollapseShapePatterns(Operation *rootOp,
                                             bool useForeach) { … }

namespace {
class DummyTrackingListener : public transform::TrackingListener { … };
} // namespace

static LogicalResult testTrackingListenerReplacements(Operation *rootOp) { … }

void TestTensorTransforms::runOnOperation() { … }

namespace mlir {
namespace test {
void registerTestTensorTransforms() { … }
} // namespace test
} // namespace mlir