| 29 | using PatternMatcherTest = HloTestBase; |
| 30 | |
| 31 | TEST_F(PatternMatcherTest, AddOp) { |
| 32 | constexpr char kModuleStr[] = R"(HloModule two_plus_two_module |
| 33 | ENTRY %two_plus_two_computation () -> f32[] { |
| 34 | %two = f32[] constant(2) |
| 35 | ROOT %two_plus_two = f32[] add(f32[] %two, f32[] %two) |
| 36 | } |
| 37 | )"; |
| 38 | TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, |
| 39 | ParseAndReturnVerifiedModule(kModuleStr)); |
| 40 | |
| 41 | const HloInstruction* matched_inst; |
| 42 | HloInstruction* matched_operand; |
| 43 | Shape* matched_shape; |
| 44 | Layout* matched_layout; |
| 45 | |
| 46 | ASSERT_TRUE(Match( |
| 47 | hlo_module->entry_computation()->root_instruction(), |
| 48 | match::Op(&matched_inst) |
| 49 | .WithName("two_plus_two") |
| 50 | .WithOpcode(HloOpcode::kAdd) |
| 51 | .WithShape( |
| 52 | match::Shape(&matched_shape) |
| 53 | .WithLayout(match::Layout(&matched_layout).WithDenseFormat())) |
| 54 | .WithOperand( |
| 55 | 0, |
| 56 | match::Op(&matched_operand).WithOpcode(HloOpcode::kConstant)))); |
| 57 | ASSERT_NE(matched_inst, nullptr); |
| 58 | EXPECT_EQ(matched_inst->name(), "two_plus_two"); |
| 59 | EXPECT_EQ(matched_inst->opcode(), HloOpcode::kAdd); |
| 60 | |
| 61 | EXPECT_TRUE(Match(hlo_module->entry_computation()->root_instruction(), |
| 62 | match::Add(match::Constant(), match::Constant()))); |
| 63 | |
| 64 | EXPECT_FALSE(Match(hlo_module->entry_computation()->root_instruction(), |
| 65 | match::Op().WithName("bad_name"))); |
| 66 | matched_inst = nullptr; |
| 67 | EXPECT_FALSE(Match(hlo_module->entry_computation()->root_instruction(), |
| 68 | match::Multiply(&matched_inst, match::Op(), match::Op()))); |
| 69 | } |
| 70 | |
| 71 | TEST_F(PatternMatcherTest, ScalarShape) { |
| 72 | auto scalar_shape = ShapeUtil::MakeShape(F32, {}); |
nothing calls this directly
no test coverage detected