feat: elide attribute type with new TypeInferableAttrInterface

This allows to omit the type of an attribute from the assembly if it can be inferred from the attribute value. For example, the type is redundant in `#substrait.timestamp<100us> : !substrait.timestamp`. The built-in `TypedAttrInterface`, however, forces the appearance of the type in the assembly. The new interface is almost identical but does not enforce it. The PR also makes the two timestamp attributes implement that interface. Signed-off-by: Ingo Müller <[email protected]>
substrait-io · Jan 21, 2025 · 1877e60 · 1877e60
1 parent c01f183
commit 1877e60
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 24 deletions.
diff --git a/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt b/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt
@@ -16,6 +16,8 @@ add_public_tablegen_target(MLIRSubstraitAttrsIncGen)
 add_dependencies(MLIRSubstraitDialect MLIRSubstraitAttrsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS SubstraitInterfaces.td)
+mlir_tablegen(SubstraitAttrInterfaces.h.inc -gen-attr-interface-decls)
+mlir_tablegen(SubstraitAttrInterfaces.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(SubstraitOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(SubstraitOpInterfaces.cpp.inc -gen-op-interface-defs)
 mlir_tablegen(SubstraitTypeInterfaces.h.inc -gen-type-interface-decls)

diff --git a/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h b/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h
@@ -19,6 +19,7 @@
 
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsDialect.h.inc" // IWYU: export
 
+#include "substrait-mlir/Dialect/Substrait/IR/SubstraitAttrInterfaces.h.inc" // IWYU: export
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpInterfaces.h.inc" // IWYU: export
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitTypeInterfaces.h.inc" // IWYU: export
 

diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitInterfaces.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitInterfaces.td
@@ -11,6 +11,26 @@
 
 include "mlir/IR/OpBase.td"
 
+def TypeInferableAttrInterface : AttrInterface<"TypeInferableAttrInterface"> {
+  let cppNamespace = "::mlir::substrait";
+  let description = [{
+    This interface is used for attributes that have a type that can be inferred
+    from the instance of the attribute. It is similar to the built-in
+    `TypedAttrInterface` in that that type is understood to represent the type
+    of the data contained in the attribute. However, it is different in that
+    `TypedAttrInterface` is typically used for cases where the type is a
+    parameter of the attribute such that there can be attribute instances with
+    the same value but different types. With this interface, the type must be
+    inferable from the value such that two instances with the same value always
+    have the same type. Crucially, this allows to elide the type in the assembly
+    format of the attribute.
+  }];
+  let methods = [InterfaceMethod<
+    "Get the attribute's type",
+    "::mlir::Type", "getType"
+  >];
+}
+
 def Substrait_ExpressionOpInterface : OpInterface<"ExpressionOpInterface"> {
   let description = [{
     Interface for any expression in a Substrait plan. This corresponds to an

diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
@@ -10,6 +10,7 @@
 #define SUBSTRAIT_DIALECT_SUBSTRAIT_IR_SUBSTRAITTYPES
 
 include "substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td"
+include "substrait-mlir/Dialect/Substrait/IR/SubstraitInterfaces.td"
 include "mlir/IR/CommonTypeConstraints.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
@@ -21,9 +22,24 @@ class Substrait_Type<string name, string typeMnemonic, list<Trait> traits = []>
 }
 
 // Base class for Substrait dialect attribute types.
-class Substrait_Attr<string name, string typeMnemonic, list<Trait> traits = []>
+class Substrait_Attr<string name, string attrMnemonic, list<Trait> traits = []>
     : AttrDef<Substrait_Dialect, name, traits> {
-  let mnemonic = typeMnemonic;
+  let mnemonic = attrMnemonic;
+}
+
+// Base class for Substrait dialect attribute types that have a statically known
+// value type.
+class Substrait_StaticallyTypedAttr<
+      string name, string attrMnemonic, string typeName, list<Trait> traits = []>
+    : Substrait_Attr<name, attrMnemonic, traits # [
+      DeclareAttrInterfaceMethods<TypeInferableAttrInterface>
+    ]> {
+  let extraClassDeclaration = [{
+    /// Implement TypeInferableAttrInterface.
+    ::mlir::Type getType() {
+      return ::mlir::substrait::}] # typeName # [{::get(getContext());
+    }
+  }];
 }
 
 def Substrait_StringType : Substrait_Type<"String", "string"> {
@@ -47,19 +63,14 @@ def Substrait_TimestampType : Substrait_Type<"Timestamp", "timestamp"> {
   }];
 }
 
-def Substrait_TimestampAttr : Substrait_Attr<"Timestamp", "timestamp",
-  [TypedAttrInterface]> {
+def Substrait_TimestampAttr
+  : Substrait_StaticallyTypedAttr<"Timestamp", "timestamp", "TimestampType"> {
   let summary = "Substrait timezone-unaware timestamp type";
   let description = [{
     This type represents a substrait timezone-unaware timestamp attribute type.
   }];  
   let parameters = (ins "int64_t":$value);
   let assemblyFormat = [{ `<` $value `` `us` `>` }];
-  let extraClassDeclaration = [{ 
-    ::mlir::Type getType() const {
-      return TimestampType::get(getContext());
-    }
-  }];
 }
 
 def Substrait_TimestampTzType : Substrait_Type<"TimestampTz", "timestamp_tz"> {
@@ -69,19 +80,14 @@ def Substrait_TimestampTzType : Substrait_Type<"TimestampTz", "timestamp_tz"> {
   }];
 }
 
-def Substrait_TimestampTzAttr : Substrait_Attr<"TimestampTz", "timestamp_tz",
-  [TypedAttrInterface]> {
+def Substrait_TimestampTzAttr : Substrait_StaticallyTypedAttr<
+    "TimestampTz", "timestamp_tz", "TimestampTzType"> {
   let summary = "Substrait timezone-aware timestamp type";
   let description = [{
     This type represents a substrait timezone-aware timestamp attribute type.
   }];  
   let parameters = (ins "int64_t":$value);
   let assemblyFormat = [{ `<` $value `` `us` `>` }];
-  let extraClassDeclaration = [{ 
-    ::mlir::Type getType() const {
-      return TimestampTzType::get(getContext());
-    }
-  }];
 }
 
 /// Currently supported atomic types. These correspond directly to the types in

diff --git a/lib/Dialect/Substrait/IR/Substrait.cpp b/lib/Dialect/Substrait/IR/Substrait.cpp
@@ -62,6 +62,7 @@ LogicalResult AdvancedExtensionAttr::verify(
 // Substrait interfaces
 //===----------------------------------------------------------------------===//
 
+#include "substrait-mlir/Dialect/Substrait/IR/SubstraitAttrInterfaces.cpp.inc"
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpInterfaces.cpp.inc"
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitTypeInterfaces.cpp.inc"
 
@@ -297,15 +298,19 @@ LiteralOp::inferReturnTypes(MLIRContext *context, std::optional<Location> loc,
                             OpaqueProperties properties, RegionRange regions,
                             llvm::SmallVectorImpl<Type> &inferredReturnTypes) {
   auto *typedProperties = properties.as<Properties *>();
+  Attribute valueAttr = typedProperties->getValue();
 
-  auto attr = llvm::dyn_cast<TypedAttr>(typedProperties->getValue());
-  if (!attr)
-    return emitOptionalError(loc, "unsuited attribute for literal value: ",
-                             typedProperties->getValue());
+  Type resultType;
+  if (auto attr = llvm::dyn_cast<TypedAttr>(valueAttr))
+    resultType = attr.getType();
+  if (auto attr = llvm::dyn_cast<TypeInferableAttrInterface>(valueAttr))
+    resultType = attr.getType();
 
-  Type resultType = attr.getType();
-  inferredReturnTypes.emplace_back(resultType);
+  if (!resultType)
+    emitOptionalError(loc, "unsuited attribute for literal value: ",
+                      typedProperties->getValue());
 
+  inferredReturnTypes.emplace_back(resultType);
   return success();
 }
 

diff --git a/test/Dialect/Substrait/literal.mlir b/test/Dialect/Substrait/literal.mlir
@@ -6,8 +6,8 @@
 // CHECK:         %[[V0:.*]] = named_table
 // CHECK-NEXT:    %[[V1:.*]] = project %[[V0]] : tuple<si1> -> tuple<si1, !substrait.timestamp, !substrait.timestamp_tz> {
 // CHECK-NEXT:    ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
-// CHECK-NEXT:      %[[V2:.*]] = literal #substrait.timestamp<10000000000us> 
-// CHECK-NEXT:      %[[V3:.*]] = literal #substrait.timestamp_tz<10000000000us> 
+// CHECK-NEXT:      %[[V2:.*]] = literal #substrait.timestamp<10000000000us>{{$}}
+// CHECK-NEXT:      %[[V3:.*]] = literal #substrait.timestamp_tz<10000000000us>{{$}}
 // CHECK-NEXT:      yield %[[V2]], %[[V3]] : !substrait.timestamp, !substrait.timestamp_tz
 // CHECK-NEXT:    }
 // CHECK-NEXT:    yield %[[V1]] : tuple<si1, !substrait.timestamp, !substrait.timestamp_tz>