198 lines · cpp
1//===- LoopLikeSCFOpsTest.cpp - SCF LoopLikeOpInterface Tests -------------===//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/Affine/IR/AffineOps.h"10#include "mlir/Dialect/Arith/IR/Arith.h"11#include "mlir/Dialect/SCF/IR/SCF.h"12#include "mlir/Dialect/SCF/Utils/Utils.h"13#include "mlir/Dialect/Utils/StaticValueUtils.h"14#include "mlir/IR/Diagnostics.h"15#include "mlir/IR/MLIRContext.h"16#include "mlir/IR/OwningOpRef.h"17#include "mlir/IR/PatternMatch.h"18#include "mlir/Interfaces/LoopLikeInterface.h"19#include "gtest/gtest.h"20 21using namespace mlir;22using namespace mlir::scf;23 24//===----------------------------------------------------------------------===//25// Test Fixture26//===----------------------------------------------------------------------===//27 28class SCFLoopLikeTest : public ::testing::Test {29protected:30 SCFLoopLikeTest() : b(&context), loc(UnknownLoc::get(&context)) {31 context.loadDialect<affine::AffineDialect, arith::ArithDialect,32 scf::SCFDialect>();33 }34 35 void checkUnidimensional(LoopLikeOpInterface loopLikeOp) {36 std::optional<OpFoldResult> maybeSingleLb =37 loopLikeOp.getSingleLowerBound();38 EXPECT_TRUE(maybeSingleLb.has_value());39 std::optional<OpFoldResult> maybeSingleUb =40 loopLikeOp.getSingleUpperBound();41 EXPECT_TRUE(maybeSingleUb.has_value());42 std::optional<OpFoldResult> maybeSingleStep = loopLikeOp.getSingleStep();43 EXPECT_TRUE(maybeSingleStep.has_value());44 std::optional<OpFoldResult> maybeSingleIndVar =45 loopLikeOp.getSingleInductionVar();46 EXPECT_TRUE(maybeSingleIndVar.has_value());47 48 std::optional<SmallVector<OpFoldResult>> maybeLb =49 loopLikeOp.getLoopLowerBounds();50 ASSERT_TRUE(maybeLb.has_value());51 EXPECT_EQ((*maybeLb).size(), 1u);52 std::optional<SmallVector<OpFoldResult>> maybeUb =53 loopLikeOp.getLoopUpperBounds();54 ASSERT_TRUE(maybeUb.has_value());55 EXPECT_EQ((*maybeUb).size(), 1u);56 std::optional<SmallVector<OpFoldResult>> maybeStep =57 loopLikeOp.getLoopSteps();58 ASSERT_TRUE(maybeStep.has_value());59 EXPECT_EQ((*maybeStep).size(), 1u);60 std::optional<SmallVector<Value>> maybeInductionVars =61 loopLikeOp.getLoopInductionVars();62 ASSERT_TRUE(maybeInductionVars.has_value());63 EXPECT_EQ((*maybeInductionVars).size(), 1u);64 }65 66 void checkMultidimensional(LoopLikeOpInterface loopLikeOp) {67 std::optional<OpFoldResult> maybeSingleLb =68 loopLikeOp.getSingleLowerBound();69 EXPECT_FALSE(maybeSingleLb.has_value());70 std::optional<OpFoldResult> maybeSingleUb =71 loopLikeOp.getSingleUpperBound();72 EXPECT_FALSE(maybeSingleUb.has_value());73 std::optional<OpFoldResult> maybeSingleStep = loopLikeOp.getSingleStep();74 EXPECT_FALSE(maybeSingleStep.has_value());75 std::optional<OpFoldResult> maybeSingleIndVar =76 loopLikeOp.getSingleInductionVar();77 EXPECT_FALSE(maybeSingleIndVar.has_value());78 79 std::optional<SmallVector<OpFoldResult>> maybeLb =80 loopLikeOp.getLoopLowerBounds();81 ASSERT_TRUE(maybeLb.has_value());82 EXPECT_EQ((*maybeLb).size(), 2u);83 std::optional<SmallVector<OpFoldResult>> maybeUb =84 loopLikeOp.getLoopUpperBounds();85 ASSERT_TRUE(maybeUb.has_value());86 EXPECT_EQ((*maybeUb).size(), 2u);87 std::optional<SmallVector<OpFoldResult>> maybeStep =88 loopLikeOp.getLoopSteps();89 ASSERT_TRUE(maybeStep.has_value());90 EXPECT_EQ((*maybeStep).size(), 2u);91 std::optional<SmallVector<Value>> maybeInductionVars =92 loopLikeOp.getLoopInductionVars();93 ASSERT_TRUE(maybeInductionVars.has_value());94 EXPECT_EQ((*maybeInductionVars).size(), 2u);95 }96 97 void checkNormalized(LoopLikeOpInterface loopLikeOp) {98 std::optional<SmallVector<OpFoldResult>> maybeLb =99 loopLikeOp.getLoopLowerBounds();100 ASSERT_TRUE(maybeLb.has_value());101 std::optional<SmallVector<OpFoldResult>> maybeStep =102 loopLikeOp.getLoopSteps();103 ASSERT_TRUE(maybeStep.has_value());104 105 auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {106 return llvm::all_of(results, [&](OpFoldResult ofr) {107 auto intValue = getConstantIntValue(ofr);108 return intValue.has_value() && intValue == val;109 });110 };111 EXPECT_TRUE(allEqual(*maybeLb, 0));112 EXPECT_TRUE(allEqual(*maybeStep, 1));113 }114 115 MLIRContext context;116 OpBuilder b;117 Location loc;118};119 120TEST_F(SCFLoopLikeTest, queryUnidimensionalLooplikes) {121 OwningOpRef<arith::ConstantIndexOp> lb =122 arith::ConstantIndexOp::create(b, loc, 0);123 OwningOpRef<arith::ConstantIndexOp> ub =124 arith::ConstantIndexOp::create(b, loc, 10);125 OwningOpRef<arith::ConstantIndexOp> step =126 arith::ConstantIndexOp::create(b, loc, 2);127 128 OwningOpRef<scf::ForOp> forOp =129 scf::ForOp::create(b, loc, lb.get(), ub.get(), step.get());130 checkUnidimensional(forOp.get());131 132 OwningOpRef<scf::ForallOp> forallOp = scf::ForallOp::create(133 b, loc, ArrayRef<OpFoldResult>(lb->getResult()),134 ArrayRef<OpFoldResult>(ub->getResult()),135 ArrayRef<OpFoldResult>(step->getResult()), ValueRange(), std::nullopt);136 checkUnidimensional(forallOp.get());137 138 OwningOpRef<scf::ParallelOp> parallelOp = scf::ParallelOp::create(139 b, loc, ValueRange(lb->getResult()), ValueRange(ub->getResult()),140 ValueRange(step->getResult()), ValueRange());141 checkUnidimensional(parallelOp.get());142}143 144TEST_F(SCFLoopLikeTest, queryMultidimensionalLooplikes) {145 OwningOpRef<arith::ConstantIndexOp> lb =146 arith::ConstantIndexOp::create(b, loc, 0);147 OwningOpRef<arith::ConstantIndexOp> ub =148 arith::ConstantIndexOp::create(b, loc, 10);149 OwningOpRef<arith::ConstantIndexOp> step =150 arith::ConstantIndexOp::create(b, loc, 2);151 152 OwningOpRef<scf::ForallOp> forallOp = scf::ForallOp::create(153 b, loc, ArrayRef<OpFoldResult>({lb->getResult(), lb->getResult()}),154 ArrayRef<OpFoldResult>({ub->getResult(), ub->getResult()}),155 ArrayRef<OpFoldResult>({step->getResult(), step->getResult()}),156 ValueRange(), std::nullopt);157 checkMultidimensional(forallOp.get());158 159 OwningOpRef<scf::ParallelOp> parallelOp = scf::ParallelOp::create(160 b, loc, ValueRange({lb->getResult(), lb->getResult()}),161 ValueRange({ub->getResult(), ub->getResult()}),162 ValueRange({step->getResult(), step->getResult()}), ValueRange());163 checkMultidimensional(parallelOp.get());164}165 166TEST_F(SCFLoopLikeTest, testForallNormalize) {167 OwningOpRef<arith::ConstantIndexOp> lb =168 arith::ConstantIndexOp::create(b, loc, 1);169 OwningOpRef<arith::ConstantIndexOp> ub =170 arith::ConstantIndexOp::create(b, loc, 10);171 OwningOpRef<arith::ConstantIndexOp> step =172 arith::ConstantIndexOp::create(b, loc, 3);173 174 scf::ForallOp forallOp = scf::ForallOp::create(175 b, loc, ArrayRef<OpFoldResult>({lb->getResult(), lb->getResult()}),176 ArrayRef<OpFoldResult>({ub->getResult(), ub->getResult()}),177 ArrayRef<OpFoldResult>({step->getResult(), step->getResult()}),178 ValueRange(), std::nullopt);179 // Create a user of the induction variable. Bitcast is chosen for simplicity180 // since it is unary.181 b.setInsertionPointToStart(forallOp.getBody());182 arith::BitcastOp::create(b, UnknownLoc::get(&context), b.getF64Type(),183 forallOp.getInductionVar(0));184 IRRewriter rewriter(b);185 FailureOr<scf::ForallOp> maybeNormalizedForallOp =186 normalizeForallOp(rewriter, forallOp);187 EXPECT_TRUE(succeeded(maybeNormalizedForallOp));188 OwningOpRef<scf::ForallOp> normalizedForallOp(*maybeNormalizedForallOp);189 checkNormalized(normalizedForallOp.get());190 191 // Check that the IV user has been updated to use the denormalized variable.192 Block *body = normalizedForallOp->getBody();193 auto bitcastOps = body->getOps<arith::BitcastOp>();194 ASSERT_EQ(std::distance(bitcastOps.begin(), bitcastOps.end()), 1);195 arith::BitcastOp ivUser = *bitcastOps.begin();196 ASSERT_NE(ivUser.getIn(), normalizedForallOp->getInductionVar(0));197}198