[mlir] load dialects for non-namespaced attrs

The mlir-translate tool calls into the parser without loading registered
dependent dialects, and the parser only loads attributes if the
fully-namespaced attribute is present in the textual IR. This causes
parsing to break when an op has an attribute that prints/parses
without the namespaced attribute.

Author: j2kun

mlir/test/Target/LLVMIR/test.mlir

@@ -40,3 +40,18 @@ llvm.func @dialect_attr_translation_multi(%a: i64, %b: i64, %c: i64) -> i64 {
 // CHECK-DAG: ![[MD_ID_ADD]] = !{!"annotation_from_test: add"}
 // CHECK-DAG: ![[MD_ID_MUL]] = !{!"annotation_from_test: mul"}
 // CHECK-DAG: ![[MD_ID_RET]] = !{!"annotation_from_test: ret"}
+
+
+// -----
+
+// This is a regression test for a bug where, during an mlir-translate call the
+// parser would only load the dialect if the fully namespaced attribute was
+// present in the IR.
+#attr = #test.nested_polynomial<<1 + x**2>>
+llvm.func @parse_correctly() {
+  // CHECK: <1 + x**2>
+  test.containing_int_polynomial_attr #attr
+
+  return
+}
+

mlir/test/lib/Dialect/Test/CMakeLists.txt

@@ -16,8 +16,8 @@ mlir_tablegen(TestOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestOps.td)
-mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls)
-mlir_tablegen(TestAttrDefs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls -attrdefs-dialect=test)
+mlir_tablegen(TestAttrDefs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=test)
 add_public_tablegen_target(MLIRTestAttrDefIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestTypeDefs.td)
@@ -86,6 +86,7 @@ add_mlir_library(MLIRTestDialect
   MLIRLinalgTransforms
   MLIRLLVMDialect
   MLIRPass
+  MLIRPolynomialDialect
   MLIRReduce
   MLIRTensorDialect
   MLIRTransformUtils

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

@@ -16,6 +16,7 @@
 // To get the test dialect definition.
 include "TestDialect.td"
 include "TestEnumDefs.td"
+include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
 include "mlir/Dialect/Utils/StructuredOpsUtils.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
@@ -351,4 +352,12 @@ def TestCustomFloatAttr : Test_Attr<"TestCustomFloat"> {
   }];
 }
 
+def NestedPolynomialAttr : Test_Attr<"NestedPolynomialAttr"> {
+  let mnemonic = "nested_polynomial";
+  let parameters = (ins Polynomial_IntPolynomialAttr:$poly);
+  let assemblyFormat = [{
+    `<` $poly `>`
+  }];
+}
+
 #endif // TEST_ATTRDEFS

mlir/test/lib/Dialect/Test/TestAttributes.h

@@ -17,6 +17,7 @@
 #include <tuple>
 
 #include "TestTraits.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"

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

@@ -232,6 +232,11 @@ def FloatElementsAttrOp : TEST_Op<"float_elements_attr"> {
   );
 }
 
+def ContainingIntPolynomialAttrOp : TEST_Op<"containing_int_polynomial_attr"> {
+  let arguments = (ins NestedPolynomialAttr:$attr);
+  let assemblyFormat = "$attr attr-dict";
+}
+
 // A pattern that updates dense<[3.0, 4.0]> to dense<[5.0, 6.0]>.
 // This tests both matching and generating float elements attributes.
 def UpdateFloatElementsAttr : Pat<
@@ -2204,7 +2209,7 @@ def ForwardBufferOp : TEST_Op<"forward_buffer", [Pure]> {
 def ReifyBoundOp : TEST_Op<"reify_bound", [Pure]> {
   let description = [{
     Reify a bound for the given index-typed value or dimension size of a shaped
-    value. "LB", "EQ" and "UB" bounds are supported. If `scalable` is set, 
+    value. "LB", "EQ" and "UB" bounds are supported. If `scalable` is set,
     `vscale_min` and `vscale_max` must be provided, which allows computing
     a bound in terms of "vector.vscale" for a given range of vscale.
   }];

mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp

@@ -15,6 +15,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialDialect.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/LLVMTranslationInterface.h"
@@ -127,6 +128,7 @@ void registerTestToLLVMIR() {
       },
       [](DialectRegistry &registry) {
         registry.insert<test::TestDialect>();
+        registry.insert<polynomial::PolynomialDialect>();
         registerBuiltinDialectTranslation(registry);
         registerLLVMDialectTranslation(registry);
         registry.addExtension(