124 lines · cpp
1//===- TestSimplification.cpp - Test simplification -----------------------===//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#include "mlir/Dialect/Func/IR/FuncOps.h"10#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"11#include "mlir/Dialect/Shard/IR/ShardOps.h"12#include "mlir/Dialect/Shard/Transforms/Partition.h"13#include "mlir/IR/BuiltinDialect.h"14#include "mlir/IR/BuiltinOps.h"15#include "mlir/IR/BuiltinTypeInterfaces.h"16#include "mlir/IR/Diagnostics.h"17#include "mlir/IR/ImplicitLocOpBuilder.h"18#include "mlir/IR/PatternMatch.h"19#include "mlir/IR/SymbolTable.h"20#include "mlir/IR/Value.h"21#include "mlir/Pass/Pass.h"22#include "mlir/Transforms/GreedyPatternRewriteDriver.h"23 24using namespace mlir;25using namespace mlir::shard;26 27namespace {28 29struct TestReshardingRewritePattern : OpRewritePattern<ShardOp> {30 using OpRewritePattern<ShardOp>::OpRewritePattern;31 32 LogicalResult matchAndRewrite(ShardOp op,33 PatternRewriter &rewriter) const override {34 if (op.getAnnotateForUsers()) {35 return failure();36 }37 38 SymbolTableCollection symbolTable;39 shard::GridOp grid = symbolTable.lookupNearestSymbolFrom<shard::GridOp>(40 op, cast<ShardingOp>(op.getSharding().getDefiningOp()).getGridAttr());41 42 bool foundUser = false;43 for (auto user : op->getUsers()) {44 if (auto targetShardOp = llvm::dyn_cast<ShardOp>(user)) {45 if (targetShardOp.getAnnotateForUsers() &&46 grid == symbolTable.lookupNearestSymbolFrom<shard::GridOp>(47 targetShardOp,48 cast<ShardingOp>(49 targetShardOp.getSharding().getDefiningOp())50 .getGridAttr())) {51 foundUser = true;52 break;53 }54 }55 }56 57 if (!foundUser) {58 return failure();59 }60 61 for (auto user : op->getUsers()) {62 auto targetShardOp = llvm::dyn_cast<ShardOp>(user);63 if (!targetShardOp || !targetShardOp.getAnnotateForUsers() ||64 symbolTable.lookupNearestSymbolFrom<shard::GridOp>(65 targetShardOp,66 cast<ShardingOp>(targetShardOp.getSharding().getDefiningOp())67 .getGridAttr()) != grid) {68 continue;69 }70 71 ImplicitLocOpBuilder builder(op->getLoc(), rewriter);72 ShapedType sourceShardShape =73 shardShapedType(op.getResult().getType(), grid, op.getSharding());74 TypedValue<ShapedType> sourceShard = cast<TypedValue<ShapedType>>(75 UnrealizedConversionCastOp::create(builder, sourceShardShape,76 op.getSrc())77 ->getResult(0));78 TypedValue<ShapedType> targetShard =79 reshard(builder, grid, op, targetShardOp, sourceShard);80 Value newTargetUnsharded =81 UnrealizedConversionCastOp::create(82 builder, targetShardOp.getResult().getType(), targetShard)83 ->getResult(0);84 rewriter.replaceAllUsesWith(targetShardOp.getResult(),85 newTargetUnsharded);86 }87 88 return success();89 }90};91 92struct TestReshardingPass93 : public PassWrapper<TestReshardingPass, OperationPass<ModuleOp>> {94 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReshardingPass)95 96 void runOnOperation() override {97 RewritePatternSet patterns(&getContext());98 patterns.insert<TestReshardingRewritePattern>(&getContext());99 if (failed(applyPatternsGreedily(getOperation().getOperation(),100 std::move(patterns)))) {101 return signalPassFailure();102 }103 }104 void getDependentDialects(DialectRegistry ®istry) const override {105 reshardingRegisterDependentDialects(registry);106 registry.insert<BuiltinDialect>();107 }108 StringRef getArgument() const final {109 return "test-grid-resharding-partition";110 }111 StringRef getDescription() const final {112 return "Test Shard dialect resharding partition.";113 }114};115} // namespace116 117namespace mlir {118namespace test {119void registerTestReshardingPartitionPass() {120 PassRegistration<TestReshardingPass>();121}122} // namespace test123} // namespace mlir124