[MLIR] Add unit test for tblgen Op build methods

[jurahul]

Summary

• Initiate the unit test with a test that tests various variants of build() methods generated for ops with variadic operands and results.
• The intent is to migrate unit .td tests in mlir/test/mlir-tblgen that check for generated C++ code to these unit tests which test both that the generated code compiles and also is functionally correct.

mlir/test/lib/Dialect/Test/TestOps.td

let extraClassDeclaration = [{
     }
     Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
     OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test TableGen generated build() methods
+//===----------------------------------------------------------------------===//
+
+def TableGenConstant : TEST_Op<"tblgen_constant"> {
+  let results = (outs AnyType);
+}
+
+// No variadic args or results.
+def TableGenBuildOp0 : TEST_Op<"tblgen_build_0"> {
+  let arguments = (ins AnyType:$value);
+  let results = (outs AnyType:$result);
+}
+
+// Sigle variadic arg and single variadic results.
+def TableGenBuildOp1 : TEST_Op<"tblgen_build_1"> {
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let results = (outs Variadic<AnyType>:$results);
+}
+
+// Single variadic arg and non-variadic results.
+def TableGenBuildOp2 : TEST_Op<"tblgen_build_2"> {
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let results = (outs AnyType:$result);
+}
+
+// Single variadic arg and multiple variadic results.
+def TableGenBuildOp3 : TEST_Op<"tblgen_build_3", [SameVariadicResultSize]> {
+  let arguments = (ins Variadic<AnyType>:$inputs);
+  let results = (outs Variadic<AnyType>:$resultA, Variadic<AnyType>:$resultB);
+}
+
 #endif // TEST_OPS

mlir/unittests/TableGen/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS enums.td)
 mlir_tablegen(EnumsGenTest.h.inc -gen-enum-decls)
 mlir_tablegen(EnumsGenTest.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRTableGenEnumsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS structs.td)
 mlir_tablegen(StructAttrGenTest.h.inc -gen-struct-attr-decls)
 mlir_tablegen(StructAttrGenTest.cpp.inc -gen-struct-attr-defs)
 add_public_tablegen_target(MLIRTableGenStructAttrIncGen)
 
 add_mlir_unittest(MLIRTableGenTests
   EnumsGenTest.cpp
   StructsGenTest.cpp
   FormatTest.cpp
+  OpBuildGen.cpp
 )
 
 add_dependencies(MLIRTableGenTests MLIRTableGenEnumsIncGen)
 add_dependencies(MLIRTableGenTests MLIRTableGenStructAttrIncGen)
+add_dependencies(MLIRTableGenTests MLIRTestDialect)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../test/lib/Dialect/Test)
+include_directories(${CMAKE_CURRENT_BINARY_DIR}/../../test/lib/Dialect/Test)
 
 target_link_libraries(MLIRTableGenTests
-  PRIVATE MLIRTableGen MLIRIR)
+  PRIVATE MLIRTableGen MLIRIR
+  PUBLIC MLIRTestDialect
+)

mlir/unittests/TableGen/OpBuildGen.cpp

+//===- OpBuildGen.cpp - TableGen OpBuildGen Tests -------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Test TableGen generated build() methods on Operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Identifier.h"
+#include "mlir/IR/StandardTypes.h"
+#include "gmock/gmock.h"
+#include <vector>
+
+namespace mlir {
+
+//===----------------------------------------------------------------------===//
+// Test Fixture
+//===----------------------------------------------------------------------===//
+
+/// Test fixture for providing basic utilities for testing.
+class OpBuildGenTest : public ::testing::Test {
+protected:
+  OpBuildGenTest()
+      : ctx{}, builder(&ctx), loc(builder.getUnknownLoc()),
+        i32Ty(builder.getI32Type()), f32Ty(builder.getF32Type()),
+        cstI32(builder.create<TableGenConstant>(loc, i32Ty)),
+        cstF32(builder.create<TableGenConstant>(loc, f32Ty)),
+        noAttrs(), attrStorage{builder.getNamedAttr("attr0",
+                                                    builder.getBoolAttr(true)),
+                               builder.getNamedAttr(
+                                   "attr1", builder.getI32IntegerAttr(33))},
+        attrs(attrStorage) {}
+
+  // Verify that `op` has the given set of result types, operands, and
+  // attributes.
+  template <typename OpTy>
+  void VerifyOp(OpTy &&concrete_op, std::vector<Type> resultTypes,
+                std::vector<Value> operands,
+                std::vector<NamedAttribute> attrs) {
+    ASSERT_NE(concrete_op, nullptr);
+    Operation *op = concrete_op.getOperation();
+
+    EXPECT_EQ(op->getNumResults(), resultTypes.size());
+    for (unsigned idx : llvm::seq(0U, op->getNumResults()))
+      EXPECT_EQ(op->getResult(idx).getType(), resultTypes[idx]);
+
+    EXPECT_EQ(op->getNumOperands(), operands.size());
+    for (unsigned idx : llvm::seq(0U, op->getNumOperands()))
+      EXPECT_EQ(op->getOperand(idx), operands[idx]);
+
+    EXPECT_EQ(op->getAttrs().size(), attrs.size());
+    for (unsigned idx : llvm::seq<unsigned>(0U, attrs.size()))
+      EXPECT_EQ(op->getAttr(attrs[idx].first.strref()), attrs[idx].second);
+
+    concrete_op.erase();
+  }
+
+protected:
+  MLIRContext ctx;
+  OpBuilder builder;
+  Location loc;
+  Type i32Ty;
+  Type f32Ty;
+  TableGenConstant cstI32;
+  TableGenConstant cstF32;
+
+  ArrayRef<NamedAttribute> noAttrs;
+  std::vector<NamedAttribute> attrStorage;
+  ArrayRef<NamedAttribute> attrs;
+};
+
+/// Test basic build methods.
+TEST_F(OpBuildGenTest, BasicBuildMethods) {
+  // Test separate args, separate results build method.
+  auto op = builder.create<TableGenBuildOp0>(loc, i32Ty, cstI32);
+  VerifyOp(op, {i32Ty}, {cstI32}, noAttrs);
+
+  // Test separate args, collective results build method.
+  op = builder.create<TableGenBuildOp0>(loc, ArrayRef<Type>{i32Ty}, cstI32);
+  VerifyOp(op, {i32Ty}, {cstI32}, noAttrs);
+
+  // Test collective args, collective params build method.
+  // TODO: Eliminate noAttrs when default value ODS change gets in.
+  op = builder.create<TableGenBuildOp0>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32}, noAttrs);
+  VerifyOp(op, {i32Ty}, {cstI32}, noAttrs);
+
+  // Test collective args, collective results, non-empty attributes
+  op = builder.create<TableGenBuildOp0>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32}, attrs);
+  VerifyOp(op, {i32Ty}, {cstI32}, attrs);
+}
+
+/// Test build methods for an Op with a single varadic arg and a single
+/// variadic result.
+TEST_F(OpBuildGenTest, BuildMethodsSingleVariadicArgAndResult) {
+  // Test collective args, collective results method, building a unary op.
+  // TODO: eliminate empty attributes
+  auto op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty},
+                                             ArrayRef<Value>{cstI32});
+  VerifyOp(std::move(op), {i32Ty}, {cstI32}, noAttrs);
+
+  // Test collective args, collective results method, building a unary op with
+  // named attributes.
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32}, attrs);
+  VerifyOp(std::move(op), {i32Ty}, {cstI32}, attrs);
+
+  // Test collective args, collective results method, building a binary op.
+  // TODO: eliminate empty attributes
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty, f32Ty},
+                                        ArrayRef<Value>{cstI32, cstF32});
+  VerifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32, cstF32}, noAttrs);
+
+  // Test collective args, collective results method, building a binary op with
+  // named attributes.
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty, f32Ty},
+                                        ArrayRef<Value>{cstI32, cstF32}, attrs);
+  VerifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32, cstF32}, attrs);
+}
+
+/// Test build methods for an Op with a single varadic arg and a non-variadic
+/// result.
+TEST_F(OpBuildGenTest, BuildMethodsSingleVariadicArgNonVariadicResults) {
+  // Test separate arg, separate param build method.
+  auto op =
+      builder.create<TableGenBuildOp1>(loc, i32Ty, ArrayRef<Value>{cstI32});
+  VerifyOp(std::move(op), {i32Ty}, {cstI32}, noAttrs);
+
+  // Test collective params build method, no attributes.
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32});
+  VerifyOp(std::move(op), {i32Ty}, {cstI32}, noAttrs);
+
+  // Test collective params build method no attributes, 2 inputs.
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32, cstF32});
+  VerifyOp(std::move(op), {i32Ty}, {cstI32, cstF32}, noAttrs);
+
+  // Test collective params build method, non-empty attributes.
+  op = builder.create<TableGenBuildOp1>(loc, ArrayRef<Type>{i32Ty},
+                                        ArrayRef<Value>{cstI32, cstF32}, attrs);
+  VerifyOp(std::move(op), {i32Ty}, {cstI32, cstF32}, attrs);
+}
+
+/// Test build methods for an Op with a single varadic arg and multiple variadic
+/// result.
+TEST_F(OpBuildGenTest,
+       BuildMethodsSingleVariadicArgAndMultipleVariadicResults) {
+  // Test separate arg, separate param build method.
+  auto op = builder.create<TableGenBuildOp3>(loc, ArrayRef<Type>{i32Ty},
+                                             ArrayRef<Type>{f32Ty},
+                                             ArrayRef<Value>{cstI32});
+  VerifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32}, noAttrs);
+
+  // Test collective params build method, no attributes.
+  op = builder.create<TableGenBuildOp3>(loc, ArrayRef<Type>{i32Ty, f32Ty},
+                                        ArrayRef<Value>{cstI32});
+  VerifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32}, noAttrs);
+
+  // Test collective params build method, with attributes.
+  op = builder.create<TableGenBuildOp3>(loc, ArrayRef<Type>{i32Ty, f32Ty},
+                                        ArrayRef<Value>{cstI32}, attrs);
+  VerifyOp(std::move(op), {i32Ty, f32Ty}, {cstI32}, attrs);
+}
+
+} // namespace mlir