[mlir][wip] Start Shape dialect

[jpienaar]

Summary

• Add basic skeleton for Shape dialect;
• Add description of types and ops to be used;

Naming still open and pretty form would enhance this. As mentioned on
discourse, will do this piece wise for discussion/easier feedback cycles. Tried
to reduce parts while allowing enough to get idea/write tests with.

[rriddle]

mlir/include/mlir/Dialect/Shape/IR/Shape.h

4

MLIR -> LLVM

10

nit: & -> and

11

typo: relations?

[rriddle]

mlir/include/mlir/Dialect/Shape/IR/Shape.h

52

///?

56

///

69

Can we get comments on each of these types?

97

///

116

weird spacing here.

118

This header comment looks off.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

4

MLIR -> LLVM

72

Are these wrapped at 80 characters? They look like they are overflowing.

mlir/include/mlir/IR/DialectSymbolRegistry.def

27

Not this revision, but we should sort these.

mlir/lib/Dialect/Shape/DialectRegistration.cpp

4

MLIR -> LLVM

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[stellaraccident]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

89

Just "Value" now, right?

[stellaraccident]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

182

I might just be mis-understanding, but to my reading, I'm not sure this formulation is right. I believe this would yield "left aligned" broadcasting, whereas the traditional rules for a ranked broadcast are "right aligned":

OUTPUT DIM: 0 1 2 
       LHS: Z Y X
       RHS:     X

I think there is a step before where you want to align the lower ranked input dimensions to fill towards the right, with respect to the output dims. For full generality, you may also need a dimension map (similar to xhl_hlo broadcast_dimensions attributes).

[jpienaar]

Addresses review comments.

[jpienaar]

Updated, thanks

mlir/include/mlir/Dialect/Shape/IR/Shape.h

69

This made me think about having ODS generate these from the description :)

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

89

Ha, good point I had written this before that change :)

182

Good catch, I forgot to add that step. In the code I had used zip_longest and reverse iterator, but that doesn't lend itself to words that well :). Here I'm pursuing broadcasting as defined in numpy and used by ML frameworks (https://docs.scipy.org/doc/numpy/reference/ufuncs.html#broadcasting) but I could be wrong so please check.

Re: broadcast_dimensions attribute, while that is more general that is not the same semantics as what folks think of when broadcasting is mentioned. That is a combination of reshape & broadcast IMHO.

mlir/include/mlir/IR/DialectSymbolRegistry.def

27

Yeah I was very tempted ... I'll do in NFC follow up.

[jpienaar]

Ran clang-format

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/include/mlir/Dialect/CMakeLists.txt

11

Seems like this should come before SPIRV?

mlir/include/mlir/Dialect/Shape/IR/Shape.h

19

Why do you need Function.h here?

37

Any reason not to sort these?

69

Realistically, we should be able to auto-generate simple types pretty trivially.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

152

Can we sort these operations alphabetically?

[antiagainst]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

182

I'd typically think an op means some action so the name is typically a verb but here we have broadcastable which sounds like a trait. What about just name it as broadcast?

208

My 2c: join sounds better for me as it indicates some operation behind. any_of gives me the feeling that the result will be the same as one of the input, which is not always true. :)

281

Would it make sense to name it as debug_print to be more clear?

[herhut]

Great to see progress on this! Some comments.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

47

Nit: It is not used for dimensions, only, but generally as a positive integer value with support for being unknown and invalid. In that sense it is more like a size_t than a dim. Maybe shape.size?

56

I find the name type somewhat misleading. I think Shape_Shape would describe better what it is (a shape). But I get the desire to name it type. Should they all be Shape_DimType and Shape_ShapeType?

101

DimOrShape?

111

+1 to the constant op. Also, why expose a -1 here? Why not use

shape.constant unknown : !shape.size
shape.constant invalid : !shape.size
shape.constant <positive integer> : !shape.size

and then parse that internally into some special integer values? This would leave more flexibility to change the backing type to something other than int later.

135

Is there a reason for this? Also, thinking of library calls that return ranked tensors of statically unknown rank, this function would not allow me to create a shape for them. Should we also have a version that takes just the rank?

182

I agree. Broadcast seems a better name, as it broadcasts two shapes into a resulting broadcasted shape.

208

any_of is less constrained in that it does not have to pick the join of the two values, whereas join has rigid semantics. On the other hand, the description below describes the join of the two values.

222

I would expect

join ([], []) -> []
join([], [*]) -> []
join([], [?,...,?]) -> [invalid]

223

join([1,?], [2, ?, ?]) -> [invalid] ?

243

"for each every" is one too many.

264

Should this be shape.mul if it existed?

272

Does this mean the rank is unknown? Or also if parts of the shape are unknown?

If known rank was allowed, one could still implement computing the rank by counting the elements of the vector (as a somewhat useless example).

[nicolasvasilache]

I question the relevance of carrying implicit broadcasting concepts from HLO and other frontends into MLIR core.
I see this as a red flag that the required separation of concerns has not happened.
Anything related to broadcast should be on the XLA / frontend side IMO.

I do not see a mention of broadcast in the discourse discussion: https://llvm.discourse.group/t/shape-type-dialect/423/4.

[nicolasvasilache]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

191

This is frontend logic that I do not think should leak into MLIR.

191

I meant into MLIR *core*, sorry.

[stellaraccident]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

191

I tend to agree. We do need this to lower frontend dialects that choose to implement numpy-style broadcasting, but that does not mean it needs to be in core (or at least in this level of core -- I could see having a numpy_broadcast_shape op somewhere).

I suspect that there may have been some thought of needing this in some follow-on analysis algorithms. It would probably be better to discuss more in that context.

[jpienaar]

Updating to address review comments (sorting ops, types, making type names uniform, renaming broadcastable op, renaming dim type, fixing example).

[jpienaar]

Updated, thanks

In D73944#1864082, @nicolasvasilache wrote:

I question the relevance of carrying implicit broadcasting concepts from HLO and other frontends into MLIR core.

I disagree. This is about enabling frontend dialects to express their shape constraints. And need not even be implicit (e.g., explicit broadcast operator has broadcast shape op in its shape function). And I see this similar to signed integer discussion. While we have preferences and would prefer to never use implicit broadcast, many frontend dialects do and this is about enabling their modelling. Enabling frontend dialects to express their constraints is very useful. This is similar to broadcastable trait that we have at the moment.

I see this as a red flag that the required separation of concerns has not happened.
Anything related to broadcast should be on the XLA / frontend side IMO.

Are you saying we should not have any frontend dialects? This moves responsibility into op construction of the frontend time, which is often not where we wanted & optimizing frontend dialect is definitely something that we want, which might require legalizing from one frontend dialect to another.

mlir/include/mlir/Dialect/CMakeLists.txt

11

Depends on if you follow sort order case sensitive or insensitive :) I normally default to case sensitive, but can do insensitive.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

19

Ah this was left over from previous testing, thanks

37

It was a semi hierarchy in my mind, but that doesn't really matter, so sorted.

69

Exactly. Should be easy to address.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

47

Good point and I like that, renaming.

56

Yeah, Shape_Shape was just weird, changing to uniformly suffixing with Type.

111

Mostly as it made writing tests easier while allowing reusing the integer attribute parser/dense storage and so switch later. But I can start with an ArrayAttr and then do the above.

For invalid case I feel constant doesn't highlight the error state well enough and would rather that be a separate op that takes a string attribute. WDYT?

135

The one below does. The main reason was that then you need to special case value interpretation (where I was using only integers, e.g., require rank followed by dims so that you could do [?] for unranked [1, ?] for 1D ranked but unknown etc.) but with expansion to using strings, one can use * now and so rolled it into constant.

191

Correct, we need this for frontend dialects (TF, TFLite, XLA [client] HLO, etc) to be able to express their shape constraints. it need not be used by standard etc. and we could add specialization that only works for ranked types. But I see unranked inputs and broadcast behavior as common for many frontend dialects and this similar to broadcastable trait or signed integers. We can move the op into a different file.

Now, of course nobody is forced to use this and standard ops will never generate nor need to solve shape equations using this.

223

I changed notations and didn't do it correctly.

243

It was just so good it happened every every time 😁

264

D'oh yes. This was computing sum of dimensions :)

272

Unraked, updated to clarify. Correct. well and I was using this for op where one computed the product of known indices, so inside the expression it would conditionally reduce 🙂

281

Sounds good. I quite like it, made writing tests easier (although I should refactor those into non-executable tests ...).

[jpienaar]

Reduce with mul should be init with 1 instead of 0

[mehdi_amini]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

244

I think si.dim should be shape.size?

[herhut]

Thanks for starting this. I'd be happy to see this land.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

111

I like the idea. Invalid shapes should always have some reason attached to them to ease debugging whereas normal constants probably don't. Something like shape.error.

191

Ultimately, we could think of making the shape dialect and the machinery built on top of it extensible so that other front-ends can add their own high-level shape operations to model their constraints. In such a world, I would agree with @nicolasvasilache that the numpy-style broadcast operation should live with the corresponding front-end dialect.
However, as a first step, we should build a closed system that is powerful enough to model constraints for the input languages we care about (which includes HLO). This requires computing shapes for broadcast, so we should have the operation.

[jpienaar]

Fix incorrectly named si.dim to shape.size

[jpienaar]

Updated typo, thanks

[stephenneuendorffer]

It seems like this landed without enough code to compile? There's no .cpp file which includes the implementation .cpp.inc generated from tablegen?
As a result BUILD_SHARED_LIBS=on fails.

Steve

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(OpenMP)
 add_subdirectory(QuantOps)
+add_subdirectory(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(VectorOps)

mlir/include/mlir/Dialect/Shape/CMakeLists.txt

+add_subdirectory(IR)

mlir/include/mlir/Dialect/Shape/IR/CMakeLists.txt

+add_mlir_dialect(ShapeOps ShapeOps)

mlir/include/mlir/Dialect/Shape/IR/Shape.h

+//===- Shape.h - MLIR Shape dialect -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the shape dialect that is used to describe and solve shape
+// relations of MLIR operations using ShapedType.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SHAPE_IR_SHAPE_H
+#define MLIR_SHAPE_IR_SHAPE_H
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace shape {
+
+/// This dialect contains shape inference related operations and facilities.
+class ShapeDialect : public Dialect {
+public:
+  /// Create the dialect in the given `context`.
+  explicit ShapeDialect(MLIRContext *context);
+};
+
+namespace ShapeTypes {
+enum Kind {
+  Component = Type::FIRST_SHAPE_TYPE,
+  Element,
+  Shape,
+  Size,
+  ValueShape
+};
+} // namespace ShapeTypes
+
+/// The component type corresponding to shape, element type and attribute.
+class ComponentType : public Type::TypeBase<ComponentType, Type> {
+public:
+  using Base::Base;
+
+  static ComponentType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Component);
+  }
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::Component;
+  }
+};
+
+/// The element type of the shaped type.
+class ElementType : public Type::TypeBase<ElementType, Type> {
+public:
+  using Base::Base;
+
+  static ElementType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Element);
+  }
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::Element;
+  }
+};
+
+/// The shape descriptor type represents rank and dimension sizes.
+class ShapeType : public Type::TypeBase<ShapeType, Type> {
+public:
+  using Base::Base;
+
+  static ShapeType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Shape);
+  }
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) { return kind == ShapeTypes::Kind::Shape; }
+};
+
+/// The type of a single dimension.
+class SizeType : public Type::TypeBase<SizeType, Type> {
+public:
+  using Base::Base;
+
+  static SizeType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Size);
+  }
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) { return kind == ShapeTypes::Kind::Size; }
+};
+
+/// The ValueShape represents a (potentially unknown) runtime value and shape.
+class ValueShapeType : public Type::TypeBase<ValueShapeType, Type> {
+public:
+  using Base::Base;
+
+  static ValueShapeType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::ValueShape);
+  }
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::ValueShape;
+  }
+};
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Shape/IR/ShapeOps.h.inc"
+
+} // namespace shape
+} // namespace mlir
+
+#endif // MLIR_SHAPE_IR_SHAPE_H

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

+//===- Shape.td - Shape operations definition --------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the operation definition file for Shape dialect operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SHAPE_OPS
+#define SHAPE_OPS
+
+include "mlir/IR/OpBase.td"
+
+// TODO(jpienaar): Move to base.
+def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped">;
+
+//===----------------------------------------------------------------------===//
+// Shape Inference dialect definitions
+//===----------------------------------------------------------------------===//
+
+def ShapeDialect : Dialect {
+  let name = "shape";
+
+  let summary = "Types and operations for shape dialect";
+  let description = [{
+    This dialect contains operations for shape inference.
+
+    Note: Unless explicitly stated, all functions that return a shape and take
+    shapes as input, return the invalid shape if one of its operands is an
+    invalid shape. This avoids flagging multiple errors for one verification
+    failure. The dialect itself does not specify how errors should be combined
+    (there are multiple different options, from always chosing first operand,
+    concatting etc. on how to combine them).
+  }];
+
+  let cppNamespace = "shape";
+}
+
+def Shape_SizeType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<DimType>()">, "dim"> {
+  let typeDescription = [{
+    `shape.size` represents a non-negative integer with support for being
+    unknown and invalid.
+
+    Operations on `shape.size` types are specialized to handle unknown/dynamic
+    value. So, for example, `<unknown> + x == <unknown>` for all non-error `x :
+    !shape.size` (e.g., an unknown value does not become known due to addition).
+  }];
+}
+
+def Shape_ShapeType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<ShapeType>()">, "shape"> {
+  let typeDescription = [{
+    `shape.type` represents either an unranked shape, a ranked shape with
+    possibly unknown dimensions or an invalid shape. The rank is of type
+    `shape.size` and, if rank is known, the extent is a 1D tensor of type
+    `shape.size`.
+
+    Shape is printed:
+    * `[*]` if it is an unranked shape
+    * `[?, 2]` if a rank 2 tensor with one unknown dimension
+    * `[3, 4]` is a rank 2 static tensor
+    * `[]` is a scalar
+    * `[1]` is a rank 1 tensor with 1 element
+    * `[invalid]` for an invalid shape
+  }];
+}
+
+def Shape_ElementType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<ElementType>()">, "element type"> {
+  let typeDescription = [{
+    `shape.element_type` represents the element type of the ShapedType. It may
+    be unknown, error or regular element type supported by ShapedType.
+  }];
+}
+
+def Shape_ComponentType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<ComponentType>()">, "component type"> {
+  let typeDescription = [{
+    `shape.element_type` represents the element type of the ShapedType. It may
+    be unknown, error or regular element type supported by ShapedType.
+  }];
+}
+
+def Shape_ValueShapeType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<ValueShapeType>()">, "value shape"> {
+  let typeDescription = [{
+    `shape.value_shape` represents the value produced by an operation (this
+    corresponds to `Value` in the compiler) and a shape. Conceptually this is a
+    tuple of a value (potentially unknown) and `shape.type`. The value and shape
+    can either or both be unknown. If both the `value` and `shape` are known,
+    then the shape of `value` is conformant with `shape`.
+  }];
+}
+
+def Shape_ShapeOrSizeType: AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
+  "shape or size">;
+
+//===----------------------------------------------------------------------===//
+// Shape op definitions
+//===----------------------------------------------------------------------===//
+
+// Base class for the operation in this dialect
+class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
+    Op<ShapeDialect, mnemonic, traits>;
+
+def Shape_AddOp : Shape_Op<"add", [SameOperandsAndResultType]> {
+  let summary = "Addition of sizes";
+  let description = [{
+    Adds two valid sizes as follows:
+    * lhs + rhs = unknown if either lhs or rhs unknown;
+    * lhs + rhs = (int)lhs + (int)rhs if known;
+  }];
+
+  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+def Shape_BroadcastOp : Shape_Op<"broadcast", []> {
+  let summary = "Returns the broadcasted output shape of two inputs";
+  let description = [{
+    Computes the broadcasted output shape following:
+    1. If any inputs are unranked, output is unranked;
+    2. Else the input array with number of dimensions smaller than the max
+       input dimension, has 1’s prepended to its shapes and the output shape is
+       calculated as follows:
+
+           output[i] = lhs[i] if lhs[i] == rhs[i] or rhs[i] is unknown/undefined
+                     = rhs[i] if lhs[i] is unknown/undefined
+                     = lhs[i] if rhs[i] == 1
+                     = rhs[i] if lhs[i] == 1
+                     = error  if lhs[i] != rhs[i]
+
+    Op has an optional string attribute for the error case where there is no
+    broadcastable output shape possible for the given inputs.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs,
+                   OptionalAttr<StrAttr>:$error);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+def Shape_ConstantOp : Shape_Op<"constant", []> {
+  let summary = "Creates a shape constant";
+  let description = [{
+    An operation that builds a size or shape from integer or array attribute.
+    It allows for creating dynamically valued shapes by using `?` for unknown
+    values. A constant shape specified with `*` will return an unranked shape.
+
+    ```mlir
+    %x = shape.constant 10 : !shape.size
+    ```
+  }];
+
+  // TODO(jpienaar): Change to a more specialized attribute that would
+  // encapsulate the unknown parsing while using denser packing.
+  let arguments = (ins ArrayAttr:$value);
+  let results = (outs Shape_ShapeOrSizeType:$result);
+}
+
+def Shape_CreateShapeOp : Shape_Op<"create_shape", []> {
+  let summary = "Creates a shape descriptor from a tensor";
+  let description = [{
+    Creates a shape from a 1D integral tensor. The rank equals the number of
+    elements in the tensor, and extent matches the values of the elements.
+  }];
+
+  let arguments = (ins I32Tensor:$input);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+def Shape_JoinOp : Shape_Op<"join", []> {
+  let summary = "Returns the least general shape.size of its operands";
+  let description = [{
+    An operation that computes the least general shape of input operands. This
+    effectively asserts that corresponding static dimensions are equal. The
+    behavior is to match each element of the `shape.type` and propagate the most
+    restrictive information, returning an invalid shape if there are
+    contradictory requirements. E.g., using pseudo code
+
+    ```
+    shape.join([*], [*]) -> [*]
+    shape.join([*], [1, ?]) -> [1, ?]
+    shape.join([1, 2], [1, ?]) -> [1, 2]
+    shape.join([*], [1, 2]) -> [1, 2]
+    shape.join([], []) -> []
+    shape.join([], [*]) -> []
+    shape.join([], [?, ?]) -> [invalid]
+    shape.join([1, ?], [2, ?, ?]) -> [invalid]
+    ```
+
+    `shape.join` also allows specifying an optional error string, that may be
+    used to return an error to the user upon mismatch of dimensions.
+
+    ```mlir
+    %c = shape.join %a, %b, error="<reason>" : !shape.type
+    ```
+  }];
+
+  let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
+                   OptionalAttr<StrAttr>:$error);
+  let results = (outs Shape_ShapeOrSizeType:$result);
+}
+
+def Shape_MulOp : Shape_Op<"mul", [SameOperandsAndResultType]> {
+  let summary = "Multiplication of sizes";
+  let description = [{
+    Multiplies two valid sizes as follows:
+    - lhs * rhs = unknown if either lhs or rhs unknown;
+    - lhs * rhs = (int)lhs * (int)rhs if both known;
+  }];
+
+  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+def Shape_ReduceOp : Shape_Op<"reduce", []> {
+  let summary = "Returns an expression reduced over a shape";
+  let description = [{
+    An operation that takes as input a shape, number of initial values and has a
+    region/function that is applied repeatedly for every dimension of the shape.
+
+    Conceptually this op performs the following reduction:
+
+    ```
+    res[] = init;
+    for (int i = 0, e = shape.rank(); i != e; ++i) {
+      res = fn(i, shape[i], res[0], ..., res[n]);
+    }
+    ```
+
+    Where fn is provided by the user and the result of the reduce op is the
+    last computed output of the reduce function. As an example, computing the
+    number of elements
+
+    ```mlir
+    func @shape_num_elements(%shape : !shape.type) -> !shape.size {
+      %0 = "shape.constant_dim"() {value = 1 : i32} : () -> !shape.size
+      %1 = "shape.reduce"(%shape, %0) ( {
+        ^bb0(%index: i32, %dim: !si.dim, %lci: !shape.size):
+          %acc = "shape.mul"(%lci, %dim) :
+            (!shape.size, !shape.size) -> !shape.size
+          "shape.return"(%acc) : (!shape.size) -> ()
+        }) : (!shape.type, !shape.size) -> (!shape.size)
+      return %1 : !shape.size
+    }
+    ```
+
+    If the shape is unranked, then the results of the op is also unranked.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$shape, Variadic<AnyType>:$args);
+  let results = (outs Variadic<AnyType>:$result);
+
+  let regions = (region SizedRegion<1>:$body);
+}
+
+def Shape_ShapeOfOp : Shape_Op<"shape_of", []> {
+  let summary = "Returns shape of a value or shaped type operand";
+
+  let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+// TODO: Add Ops: if_static, if_ranked
+
+// For testing usage.
+def Shape_DebugPrintOp : Shape_Op<"debug_print", []> {
+  let summary = "Prints the input shape or size";
+  let description = [{
+    Prints the input dim or shape and passes through input.
+
+    Note: This is intended for testing and debugging only.
+  }];
+
+  let arguments = (ins Shape_ShapeOrSizeType:$input);
+  let results =  (outs Shape_ShapeOrSizeType:$output);
+}
+
+#endif // SHAPE_OPS

mlir/include/mlir/IR/DialectSymbolRegistry.def

 //===- DialectSymbolRegistry.def - MLIR Dialect Symbol Registry -*- C++ -*-===//
 //
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// Part of the MLIR 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file enumerates the different dialects that define custom classes
 // within the attribute or type system.
 //
 //===----------------------------------------------------------------------===//
 
 DEFINE_SYM_KIND_RANGE(STANDARD)
 DEFINE_SYM_KIND_RANGE(TENSORFLOW_CONTROL)
 DEFINE_SYM_KIND_RANGE(TENSORFLOW_EXECUTOR)
 DEFINE_SYM_KIND_RANGE(TENSORFLOW)
 DEFINE_SYM_KIND_RANGE(LLVM)
 DEFINE_SYM_KIND_RANGE(QUANTIZATION)
 DEFINE_SYM_KIND_RANGE(IREE) // IREE stands for IR Execution Engine
 DEFINE_SYM_KIND_RANGE(LINALG) // Linear Algebra Dialect
 DEFINE_SYM_KIND_RANGE(FIR) // Flang Fortran IR Dialect
 DEFINE_SYM_KIND_RANGE(OPENMP) // OpenMP IR Dialect
 DEFINE_SYM_KIND_RANGE(TOY) // Toy language (tutorial) Dialect
 DEFINE_SYM_KIND_RANGE(SPIRV) // SPIR-V dialect
 DEFINE_SYM_KIND_RANGE(XLA_HLO) // XLA HLO dialect
+DEFINE_SYM_KIND_RANGE(SHAPE) // Shape dialect
 
 // The following ranges are reserved for experimenting with MLIR dialects in a
 // private context without having to register them here.
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_0)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_1)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_2)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_3)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_4)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_5)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_6)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_7)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_8)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_9)
 
 #undef DEFINE_SYM_KIND_RANGE

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(OpenMP)
 add_subdirectory(QuantOps)
 add_subdirectory(SDBM)
+add_subdirectory(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(VectorOps)
 
 add_llvm_library(MLIRDialect
   Traits.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect
   )
 target_link_libraries(MLIRDialect MLIRIR)

mlir/lib/Dialect/Shape/CMakeLists.txt

+file(GLOB globbed *.c *.cpp)
+add_llvm_library(MLIRShape
+  ${globbed}
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
+  )
+add_dependencies(MLIRShape MLIRShapeOpsIncGen LLVMSupport)
+target_link_libraries(MLIRShape LLVMSupport)

mlir/lib/Dialect/Shape/DialectRegistration.cpp

+//===- DialectRegistration.cpp - Register shape dialect -------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Shape/IR/Shape.h"
+using namespace mlir;
+
+// Static initialization for shape dialect registration.
+static DialectRegistration<shape::ShapeDialect> ShapeOps;