[mlir] Add target for Shape dialect

[jpienaar]

Summary

BEGIN_PUBLIC
[mlir] Add target for Shape dialect

Add targets and basic printing/parsing of types in Shape dialect.
END_PUBLIC

[silvas]

Thanks, this is a good start.

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

222–223

Shape_SizeType?

mlir/lib/Dialect/Shape/IR/Shape.cpp

27–28

Can we remove this? Just use "foo.op" instead of "shape.op" during testing?

mlir/test/Dialect/Shape/ops.mlir

30

Can we use "!shape.shape" it seems a lot more readable to me.

[silvas]

LGTM after comments are addressed.

[jpienaar]

Rename shape.type to shape.shape

[jpienaar]

Updated, thanks

mlir/lib/Dialect/Shape/IR/Shape.cpp

27–28

Expanded the comment to mention it is also for prototyping while the dialect is still changing a lot.

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

-add_mlir_dialect(ShapeOps shape ShapeOps)
+set(LLVM_TARGET_DEFINITIONS ShapeOps.td)
+mlir_tablegen(ShapeOps.h.inc -gen-op-decls)
+mlir_tablegen(ShapeOps.cpp.inc -gen-op-defs)
+mlir_tablegen(ShapeOpsDialect.h.inc -gen-dialect-decls)
+add_public_tablegen_target(MLIRShapeOpsIncGen)

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

 namespace shape {
 
 namespace ShapeTypes {
 enum Kind {
   Component = Type::FIRST_SHAPE_TYPE,
   Element,
   Shape,
   Size,
-  ValueShape
+  ValueShape,
+  LAST_SHAPE_TYPE = ValueShape
 };
 } // namespace ShapeTypes
 
 /// The component type corresponding to shape, element type and attribute.
 class ComponentType : public Type::TypeBase<ComponentType, Type> {
 public:
   using Base::Base;
 

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"
+include "mlir/Interfaces/SideEffects.td"
 
 // TODO(jpienaar): Move to base.
 def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped">;
 
 //===----------------------------------------------------------------------===//
 // Shape Inference dialect definitions
 //===----------------------------------------------------------------------===//
 
let description = [{
     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"> {
+def Shape_ComponentType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<::mlir::shape::ComponentType>()">, "component type"> {
   let typeDescription = [{
-    `shape.size` represents a non-negative integer with support for being
-    unknown and invalid.
+    `shape.element_type` represents the element type of the ShapedType. It may
+    be unknown, error or regular element type supported by ShapedType.
+  }];
+}
 
-    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_ElementType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<::mlir::shape::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_ShapeType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<ShapeType>()">, "shape"> {
+    CPred<"$_self.isa<::mlir::shape::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"> {
+def Shape_SizeType : DialectType<ShapeDialect,
+    CPred<"$_self.isa<::mlir::shape::SizeType>()">, "size"> {
   let typeDescription = [{
-    `shape.element_type` represents the element type of the ShapedType. It may
-    be unknown, error or regular element type supported by ShapedType.
-  }];
-}
+    `shape.size` represents a non-negative integer with support for being
+    unknown and invalid.
 
-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.
+    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_ValueShapeType : DialectType<ShapeDialect,
-    CPred<"$_self.isa<ValueShapeType>()">, "value shape"> {
+    CPred<"$_self.isa<::mlir::shape::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_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);
+  let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
+  let results = (outs Shape_SizeType:$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
let description = [{
 
     ```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 arguments = (ins AnyAttr:$value);
   let results = (outs Shape_ShapeOrSizeType:$result);
+
+  // TODO: Move this to main so that all shape ops implement these.
+  let printer = [{ return ::print(p, *this); }];
+  let verifier = [{ return ::verify(*this); }];
+  let parser = [{ return ::parse$cppClass(parser, 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.
   }];

 
 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);
+  let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
+  let results = (outs Shape_SizeType:$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.
 
let description = [{
 
     ```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: !shape.size, %lci: !shape.size):
           %acc = "shape.mul"(%lci, %dim) :
             (!shape.size, !shape.size) -> !shape.size
-          "shape.return"(%acc) : (!shape.size) -> ()
+          shape.yield %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);
 }
 
+def Shape_YieldOp : Shape_Op<"yield", [NoSideEffect, Terminator]> {
+  let summary = "Returns the value to parent op";
+
+  let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [OpBuilder<
+    "Builder *b, OperationState &result", [{ build(b, result, llvm::None); }]
+  >];
+
+  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+}
+
 // 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/InitAllDialects.h

 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/QuantOps/QuantOps.h"
 #include "mlir/Dialect/SDBM/SDBMDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/VectorOps/VectorOps.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
 // This function should be called before creating any MLIRContext if one expect
 // all the possible dialects to be made available to the context automatically.
 inline void registerAllDialects() {
   static bool init_once = []() {
     registerDialect<AffineOpsDialect>();
     registerDialect<fxpmath::FxpMathOpsDialect>();
     registerDialect<gpu::GPUDialect>();
     registerDialect<LLVM::LLVMDialect>();
     registerDialect<linalg::LinalgDialect>();
     registerDialect<loop::LoopOpsDialect>();
     registerDialect<omp::OpenMPDialect>();
     registerDialect<quant::QuantizationDialect>();
     registerDialect<spirv::SPIRVDialect>();
     registerDialect<StandardOpsDialect>();
     registerDialect<vector::VectorOpsDialect>();
     registerDialect<NVVM::NVVMDialect>();
     registerDialect<ROCDL::ROCDLDialect>();
     registerDialect<SDBMDialect>();
+    registerDialect<shape::ShapeDialect>();
     return true;
   }();
   (void)init_once;
 }
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

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(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(VectorOps)
 
 set(LLVM_OPTIONAL_SOURCES
   Traits.cpp
 )
 

mlir/lib/Dialect/Shape/CMakeLists.txt

-file(GLOB globbed *.c *.cpp)
 add_mlir_dialect_library(MLIRShape
-  ${globbed}
+  IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
+
+  DEPENDS
+  MLIRShapeOpsIncGen
   )
-add_dependencies(MLIRShape MLIRShapeOpsIncGen LLVMSupport)
 target_link_libraries(MLIRShape
   PUBLIC
+  MLIRIR
   MLIRSideEffects
-  LLVMSupport)
+  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;

mlir/lib/Dialect/Shape/IR/Shape.cpp

+//===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
+//
+// 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"
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/StandardTypes.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::shape;
+
+ShapeDialect::ShapeDialect(MLIRContext *context)
+    : Dialect(getDialectNamespace(), context) {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
+      >();
+  addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType>();
+  // Allow unknown operations during prototyping and testing. As the dialect is
+  // still evolving it makes it simple to start with an unregistered ops and
+  // try different variants before actually defining the op.
+  allowUnknownOperations();
+}
+
+/// Parse a type registered to this dialect.
+Type ShapeDialect::parseType(DialectAsmParser &parser) const {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword))
+    return Type();
+
+  if (keyword == "component")
+    return ComponentType::get(getContext());
+  if (keyword == "element")
+    return ElementType::get(getContext());
+  if (keyword == "shape")
+    return ShapeType::get(getContext());
+  if (keyword == "size")
+    return SizeType::get(getContext());
+  if (keyword == "value_shape")
+    return ValueShapeType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown shape type: ") << keyword;
+  return Type();
+}
+
+/// Print a type registered to this dialect.
+void ShapeDialect::printType(Type type, DialectAsmPrinter &os) const {
+  switch (type.getKind()) {
+  case ShapeTypes::Component:
+    os << "component";
+    return;
+  case ShapeTypes::Element:
+    os << "element";
+    return;
+  case ShapeTypes::Size:
+    os << "size";
+    return;
+  case ShapeTypes::Shape:
+    os << "shape";
+    return;
+  case ShapeTypes::ValueShape:
+    os << "value_shape";
+    return;
+  default:
+    llvm_unreachable("unexpected 'shape' type kind");
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Constant*Op
+//===----------------------------------------------------------------------===//
+
+static void print(OpAsmPrinter &p, ConstantOp &op) {
+  p << "shape.constant ";
+  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"value"});
+
+  if (op.getAttrs().size() > 1)
+    p << ' ';
+  p.printAttributeWithoutType(op.value());
+  p << " : " << op.getType();
+}
+
+static ParseResult parseConstantOp(OpAsmParser &parser,
+                                   OperationState &result) {
+  Attribute valueAttr;
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+  Type i64Type = parser.getBuilder().getIntegerType(64);
+  if (parser.parseAttribute(valueAttr, i64Type, "value", result.attributes))
+    return failure();
+
+  Type type;
+  if (parser.parseColonType(type))
+    return failure();
+
+  // Add the attribute type to the list.
+  return parser.addTypeToList(type, result.types);
+}
+
+static LogicalResult verify(ConstantOp &op) { return success(); }
+
+namespace mlir {
+namespace shape {
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
+
+} // namespace shape
+} // namespace mlir

mlir/test/Dialect/Shape/ops.mlir

+// RUN: mlir-opt -split-input-file %s | FileCheck %s --dump-input-on-failure
+
+// CHECK-LABEL: shape_num_elements
+func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
+  %0 = shape.constant 0 : !shape.size
+  %1 = "shape.reduce"(%shape, %0) ( {
+    ^bb0(%index: i32, %dim: !shape.size, %lci: !shape.size):
+      %acc = "shape.add"(%lci, %dim) : (!shape.size, !shape.size) -> !shape.size
+      "shape.yield"(%acc) : (!shape.size) -> ()
+    }) : (!shape.shape, !shape.size) -> (!shape.size)
+  return %1 : !shape.size
+}
+
+func @test_shape_num_elements_unknown() {
+  %0 = "shape.unknown_shape"() : () -> !shape.shape
+  %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
+  %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
+  return
+}
+
+func @test_shape_num_elements_fixed() {
+  %0 = "shape.constant"() { value = [1, 57, 92] }: () -> !shape.shape
+  %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
+  %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
+  return
+}
+
+func @test_broadcastable_fixed() {
+  %0 = "shape.constant"() { value = [10, 1, 57, 92] }: () -> !shape.shape
+  %1 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %2 = "shape.broadcastable"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
+  return
+}
+
+func @test_shape_any_fixed() {
+  %0 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %1 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
+  return
+}
+
+func @test_shape_any_unknown() {
+  %0 = "shape.constant"() { value = [4, -1, 92] }: () -> !shape.shape
+  %1 = "shape.constant"() { value = [-1, 57, 92] }: () -> !shape.shape
+  %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
+  return
+}
+
+func @test_shape_any_fixed_mismatch() {
+  %0 = "shape.constant"() { value = [4, 57, 92] }: () -> !shape.shape
+  %1 = "shape.constant"() { value = [2, 57, 92] }: () -> !shape.shape
+  %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
+  return
+}