[mlir] Add a ValueSemantics trait. #99493

pifon2a · 2024-07-18T13:41:12Z

We need to distinguish ShapedTypes with and without value semantics. This is needed for downstream users to define their custom vector and tensor
types that can work with the arith/math dialect.

RFC https://discourse.llvm.org/t/rfc-mlir-types-with-encoding/80189

llvmbot · 2024-07-18T13:41:44Z

@llvm/pr-subscribers-mlir-ods
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Alexander Belyaev (pifon2a)

Changes

This is needed for downstream users to define their custom vector and tensor types that can work with the arith/math dialect.

RFC https://discourse.llvm.org/t/rfc-mlir-types-with-encoding/80189

Full diff: https://github.com/llvm/llvm-project/pull/99493.diff

3 Files Affected:

(modified) mlir/include/mlir/IR/BuiltinTypes.h (+5)
(modified) mlir/include/mlir/IR/BuiltinTypes.td (+12-2)
(modified) mlir/include/mlir/IR/CommonTypeConstraints.td (+16-4)

diff --git a/mlir/include/mlir/IR/BuiltinTypes.h b/mlir/include/mlir/IR/BuiltinTypes.h
index 5579b138668d2..ab01ea2293f3d 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.h
+++ b/mlir/include/mlir/IR/BuiltinTypes.h
@@ -39,6 +39,11 @@ struct IntegerTypeStorage;
 struct TupleTypeStorage;
 } // namespace detail
 
+/// Type trait indicating that the type can be an operand to an elementwise op.
+template <typename ConcreteType>
+class MappableContainer
+    : public TypeTrait::TraitBase<ConcreteType, MappableContainer> {};
+
 //===----------------------------------------------------------------------===//
 // FloatType
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 4cade83dd3c32..9e93593bb45d4 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -30,6 +30,15 @@ class Builtin_Type<string name, string typeMnemonic, list<Trait> traits = [],
   let typeName = "builtin." # typeMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// Traits
+//===----------------------------------------------------------------------===//
+
+/// Type trait indicating that the type can be an operand to an elementwise op.
+def MappableContainer : NativeTypeTrait<"MappableContainer"> {
+  let cppNamespace = "::mlir";
+}
+
 //===----------------------------------------------------------------------===//
 // ComplexType
 //===----------------------------------------------------------------------===//
@@ -745,7 +754,7 @@ def Builtin_Opaque : Builtin_Type<"Opaque", "opaque"> {
 //===----------------------------------------------------------------------===//
 
 def Builtin_RankedTensor : Builtin_Type<"RankedTensor", "tensor", [
-    ShapedTypeInterface
+    MappableContainer, ShapedTypeInterface
   ], "TensorType"> {
   let summary = "Multi-dimensional array with a fixed number of dimensions";
   let description = [{
@@ -1049,7 +1058,8 @@ def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", "unranked_tensor", [
 // VectorType
 //===----------------------------------------------------------------------===//
 
-def Builtin_Vector : Builtin_Type<"Vector", "vector", [ShapedTypeInterface], "Type"> {
+def Builtin_Vector : Builtin_Type<"Vector", "vector",
+    [MappableContainer, ShapedTypeInterface], "Type"> {
   let summary = "Multi-dimensional SIMD vector type";
   let description = [{
     Syntax:
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index af4f13dc09360..2284d59869a20 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/mlir/include/mlir/IR/CommonTypeConstraints.td
@@ -89,6 +89,9 @@ def HasStaticShapePred :
 // Whether a type is a TupleType.
 def IsTupleTypePred : CPred<"::llvm::isa<::mlir::TupleType>($_self)">;
 
+// Whether a type has a MappableContainer trait.
+def IsMappableContainerPred : CPred<"$_self.hasTrait<MappableContainer>()">;
+
 //===----------------------------------------------------------------------===//
 // Type definitions
 //===----------------------------------------------------------------------===//
@@ -403,6 +406,10 @@ class HasRankGreaterOrEqualPred<int rank> : And<[
     CPred<[{::llvm::cast<::mlir::ShapedType>($_self).getRank() >= }] # rank>
 ]>;
 
+// Mappable types.
+class MappableContainerOf<list<Type> allowedTypes> :
+  ShapedContainerType<allowedTypes, IsMappableContainerPred, "mapable container">;
+
 // Vector types.
 
 class VectorOf<list<Type> allowedTypes> :
@@ -842,10 +849,15 @@ class NestedTupleOf<list<Type> allowedTypes> :
 // Common type constraints
 //===----------------------------------------------------------------------===//
 // Type constraint for types that are "like" some type or set of types T, that is
-// they're either a T, a vector of Ts, or a tensor of Ts
+// they're either a T, a vector of Ts, or a tensor of Ts.
 class TypeOrContainer<Type allowedType, string name> : TypeConstraint<Or<[
-  allowedType.predicate, VectorOf<[allowedType]>.predicate,
-  TensorOf<[allowedType]>.predicate]>,
+  allowedType.predicate, MappableContainerOf<[allowedType]>.predicate]>,
+  name>;
+
+// Type constraint for types that are "like" some type or set of types T, that is
+// they're either a T or a mapable container of Ts.
+class TypeOrMappableContainer<Type allowedType, string name> : TypeConstraint<Or<[
+  allowedType.predicate, MappableContainerOf<[allowedType]>.predicate]>,
   name>;
 
 // Temporary constraint to allow gradual transition to supporting 0-D vectors.
@@ -864,7 +876,7 @@ def BoolLikeOfAnyRank : TypeOrContainerOfAnyRank<I1, "bool-like">;
 
 // Type constraint for signless-integer-like types: signless integers, indices,
 // vectors of signless integers or indices, tensors of signless integers.
-def SignlessIntegerLike : TypeOrContainer<AnySignlessIntegerOrIndex,
+def SignlessIntegerLike : TypeOrMappableContainer<AnySignlessIntegerOrIndex,
     "signless-integer-like">;
 
 def SignlessIntegerLikeOfAnyRank : TypeOrContainerOfAnyRank<

joker-eph · 2024-07-18T13:48:40Z

mlir/include/mlir/IR/BuiltinTypes.h

@@ -39,6 +39,11 @@ struct IntegerTypeStorage;
 struct TupleTypeStorage;
 } // namespace detail

+/// Type trait indicating that the type can be an operand to an elementwise op.


This description, as well as the name of the trait does not seem to be something that we couldn't apply to Memref.
However this isn't the intent right?

No, it is not the intent. How would you phrase it better? I want this trait to be for value types like vector or tensor, but not for memrefs. MemRefs are just pointers, it is not clear what we actually map.

Memref are more than pointers: they have a shape, which means you can actually iterate all the elements inside a memref. This makes it a candidate for a "map" function I believe (I understand this isn't what you want, which is why I'm picking on the name).
I think @jpienaar proposed something for this?

MappableValueBasedContainer?

Isn't this conflating two different traits? For an arith op on a shaped type you want it to have value semantics and you want the right element type. Why does that need to be expressed in a single trait?

Added the second trait.

joker-eph · 2024-07-19T11:47:51Z

mlir/include/mlir/IR/BuiltinTypes.h

+/// Type trait indicating that the type can be an operand to an elementwise op.
+template <typename ConcreteType>
+class MappableContainer
+    : public TypeTrait::TraitBase<ConcreteType, MappableContainer> {};


Do we need this right now knowing we have the ShapedTypeInterface?

We don't. Good catch, Mehdi!

ThomasRaoux · 2024-07-19T17:59:12Z

looks like there are some build errors on the bot:


In file included from /var/lib/buildkite-agent/builds/linux-56-59b8f5d88-tvbq4-1/llvm-project/github-pull-requests/flang/lib/Optimizer/CodeGen/CGOps.cpp:33:
--
  | /var/lib/buildkite-agent/builds/linux-56-59b8f5d88-tvbq4-1/llvm-project/github-pull-requests/build/tools/flang/include/flang/Optimizer/CodeGen/CGOps.cpp.inc:43:62: error: use of undeclared identifier 'ValueSemantics'; did you mean 'mlir::ValueSemantics'?
  | if (!((((type.isSignlessIntOrIndex())) \|\| (((type.hasTrait<ValueSemantics>())) && ([](::mlir::Type elementType) { return (elementType.isSignlessIntOrIndex()); }(::llvm::cast<::mlir::ShapedType>(type).getElementType())))) \|\| ((type.isSignedInteger())) \|\| ((::llvm::isa<::fir::IntegerType>(type))))) {
  | ^~~~~~~~~~~~~~
  | mlir::ValueSemantics

pifon2a · 2024-07-19T19:21:25Z

looks like there are some build errors on the bot:


In file included from /var/lib/buildkite-agent/builds/linux-56-59b8f5d88-tvbq4-1/llvm-project/github-pull-requests/flang/lib/Optimizer/CodeGen/CGOps.cpp:33:
--
  | /var/lib/buildkite-agent/builds/linux-56-59b8f5d88-tvbq4-1/llvm-project/github-pull-requests/build/tools/flang/include/flang/Optimizer/CodeGen/CGOps.cpp.inc:43:62: error: use of undeclared identifier 'ValueSemantics'; did you mean 'mlir::ValueSemantics'?
  | if (!((((type.isSignlessIntOrIndex())) \|\| (((type.hasTrait<ValueSemantics>())) && ([](::mlir::Type elementType) { return (elementType.isSignlessIntOrIndex()); }(::llvm::cast<::mlir::ShapedType>(type).getElementType())))) \|\| ((type.isSignedInteger())) \|\| ((::llvm::isa<::fir::IntegerType>(type))))) {
  | ^~~~~~~~~~~~~~
  | mlir::ValueSemantics

Should work now. Let's see.

joker-eph

LGTM but if @jpienaar has also time to review, it'd be better.

ThomasRaoux

Looks good to me as well

jpienaar

Looks good to me, I think it reads simply and meaning is clear (I'm reminded we are missing doc generation for these, unrelated of course).

We need to distinguish ShapedTypes with and without value semantics. This is needed for downstream users to define their custom vector and tensor types that can work with the arith/math dialect. RFC https://discourse.llvm.org/t/rfc-mlir-types-with-encoding/80189

stellaraccident

Belated lgtm. Thanks.

Summary: We need to distinguish ShapedTypes with and without value semantics. This is needed for downstream users to define their custom vector and tensor types that can work with the arith/math dialect. RFC https://discourse.llvm.org/t/rfc-mlir-types-with-encoding/80189 Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60251327

pifon2a requested review from ftynse and ThomasRaoux July 18, 2024 13:41

llvmbot added mlir:core MLIR Core Infrastructure mlir mlir:ods labels Jul 18, 2024

joker-eph reviewed Jul 18, 2024

View reviewed changes

pifon2a force-pushed the trait branch from 2674e4c to 151af4a Compare July 19, 2024 08:00

joker-eph reviewed Jul 19, 2024

View reviewed changes

pifon2a force-pushed the trait branch from 151af4a to 89870ac Compare July 19, 2024 11:58

pifon2a changed the title ~~[mlir] Add a MappableContainer trait.~~ [mlir] Add a ValueSemantics trait. Jul 19, 2024

pifon2a force-pushed the trait branch from 89870ac to d79e521 Compare July 19, 2024 19:20

joker-eph requested a review from jpienaar July 19, 2024 19:29

joker-eph approved these changes Jul 19, 2024

View reviewed changes

ThomasRaoux approved these changes Jul 19, 2024

View reviewed changes

jpienaar approved these changes Jul 19, 2024

View reviewed changes

pifon2a force-pushed the trait branch from d79e521 to a79af6e Compare July 21, 2024 20:48

pifon2a merged commit 568845a into llvm:main Jul 21, 2024
7 checks passed

stellaraccident reviewed Jul 21, 2024

View reviewed changes

[mlir] Add a ValueSemantics trait. #99493

[mlir] Add a ValueSemantics trait. #99493

Uh oh!

Conversation

pifon2a commented Jul 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jul 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux commented Jul 19, 2024

Uh oh!

pifon2a commented Jul 19, 2024

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux left a comment

Choose a reason for hiding this comment

Uh oh!

jpienaar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stellaraccident left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pifon2a commented Jul 18, 2024 •

edited

Loading

llvmbot commented Jul 18, 2024 •

edited

Loading