79 lines · cpp
1//===- TestProcessMultiIndexOpLowering.cpp --------------------------------===//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/Arith/IR/Arith.h"10#include "mlir/Dialect/Shard/Transforms/Transforms.h"11#include "mlir/Dialect/Utils/IndexingUtils.h"12#include "mlir/IR/SymbolTable.h"13#include "mlir/Pass/Pass.h"14#include "mlir/Transforms/GreedyPatternRewriteDriver.h"15 16using namespace mlir;17 18namespace {19 20struct TestAllSliceOpLoweringPass21 : public PassWrapper<TestAllSliceOpLoweringPass, OperationPass<>> {22 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestAllSliceOpLoweringPass)23 24 void runOnOperation() override {25 RewritePatternSet patterns(&getContext());26 SymbolTableCollection symbolTableCollection;27 shard::populateAllSliceOpLoweringPatterns(patterns, symbolTableCollection);28 LogicalResult status =29 applyPatternsGreedily(getOperation(), std::move(patterns));30 (void)status;31 assert(succeeded(status) && "applyPatternsGreedily failed.");32 }33 void getDependentDialects(DialectRegistry ®istry) const override {34 shard::registerAllSliceOpLoweringDialects(registry);35 }36 StringRef getArgument() const final {37 return "test-grid-all-slice-op-lowering";38 }39 StringRef getDescription() const final {40 return "Test lowering of all-slice.";41 }42};43 44struct TestMultiIndexOpLoweringPass45 : public PassWrapper<TestMultiIndexOpLoweringPass, OperationPass<>> {46 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestMultiIndexOpLoweringPass)47 48 void runOnOperation() override {49 RewritePatternSet patterns(&getContext());50 SymbolTableCollection symbolTableCollection;51 shard::populateProcessMultiIndexOpLoweringPatterns(patterns,52 symbolTableCollection);53 LogicalResult status =54 applyPatternsGreedily(getOperation(), std::move(patterns));55 (void)status;56 assert(succeeded(status) && "applyPatternsGreedily failed.");57 }58 void getDependentDialects(DialectRegistry ®istry) const override {59 shard::registerProcessMultiIndexOpLoweringDialects(registry);60 }61 StringRef getArgument() const final {62 return "test-grid-process-multi-index-op-lowering";63 }64 StringRef getDescription() const final {65 return "Test lowering of shard.process_multi_index op.";66 }67};68 69} // namespace70 71namespace mlir {72namespace test {73void registerTestOpLoweringPasses() {74 PassRegistration<TestAllSliceOpLoweringPass>();75 PassRegistration<TestMultiIndexOpLoweringPass>();76}77} // namespace test78} // namespace mlir79