[mlir] Custom printing/parsing for Shape::AssumingOp

[tpopp]

Summary

Additionally, this adds traits and builder methods to AssumingYieldOp
and names the input witness to the AssumingOp.

[herhut]

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

430–431

condition -> witness?

[rriddle]

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

145

nit: drop trivial braces.

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

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SHAPE_IR_SHAPE_H
 #define MLIR_SHAPE_IR_SHAPE_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace shape {
 
 namespace ShapeTypes {
 enum Kind {

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/Dialect/Shape/IR/ShapeBase.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Shape_WitnessType : DialectType<ShapeDialect,
     CPred<"$_self.isa<::mlir::shape::WitnessType>()">, "witness">,
     BuildableType<"$_builder.getType<::mlir::shape::WitnessType>()"> {
   let typeDescription = [{
     A witness is a structural device in the compiler to maintain ordering of
def Shape_AssumingOp : Shape_Op<"assuming",
   let description = [{
     Executes the region assuming all witnesses are true.
 
     "assuming" operations represent an execution order restriction to the
     compiler, information for dependent code to rely on (by assuming), and
     nothing else. They should not exist after a program is fully lowered and
     ready to execute.
   }];
-  let arguments = (ins Shape_WitnessType);
-  let regions = (region SizedRegion<1>:$thenRegion);
+  let arguments = (ins Shape_WitnessType:$witness);
+  let regions = (region SizedRegion<1>:$doRegion);
   let results = (outs Variadic<AnyType>:$results);
+
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def Shape_AssumingYieldOp : Shape_Op<"assuming_yield", [Terminator]> {
+def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
+                                     [NoSideEffect, ReturnLike, Terminator]> {
   let summary = "Yield operation";
   let description = [{
     This yield operation represents a return operation within the assert_and_exec
     region. The operation takes variable number of operands and produces no
     results. The operand number and types must match the return signature of
     the region that contains the operation.
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
+
+  let builders = [
+    OpBuilder<"OpBuilder &builder, OperationState &result",
+    [{ /* nothing to do */ }]>
+  ];
 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable", []> {
   let summary = "Determines if 2 shapes can be successfully broadcasted.";
   let description = [{
     Given 2 input shapes, return a witness specifying if they are broadcastable.
     This broadcastable follows the same logic as what shape.broadcast documents.
 

mlir/lib/Dialect/Shape/CMakeLists.txt

 add_mlir_dialect_library(MLIRShape
   IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
 
   DEPENDS
   MLIRShapeOpsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRControlFlowInterfaces
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   )

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

AnyOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                         ValueRange operands, DictionaryAttr attributes,
                         RegionRange regions,
                         SmallVectorImpl<Type> &inferredReturnTypes) {
   inferredReturnTypes.push_back(ShapeType::get(context));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// AssumingOp
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseAssumingOp(OpAsmParser &parser,
+                                   OperationState &result) {
+  result.regions.reserve(1);
+  Region *doRegion = result.addRegion();
+
+  auto &builder = parser.getBuilder();
+  OpAsmParser::OperandType cond;
+  if (parser.parseOperand(cond) ||
+      parser.resolveOperand(cond, builder.getType<WitnessType>(),
+                            result.operands))
+    return failure();
+
+  // Parse optional results type list.
+  if (parser.parseOptionalArrowTypeList(result.types))
+    return failure();
+
+  // Parse the region and add a terminator if elided.
+  if (parser.parseRegion(*doRegion, /*arguments=*/{}, /*argTypes=*/{}))
+    return failure();
+  AssumingOp::ensureTerminator(*doRegion, parser.getBuilder(), result.location);
+
+  // Parse the optional attribute list.
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+  return success();
+}
+
+static void print(OpAsmPrinter &p, AssumingOp op) {
+  bool yieldsResults = !op.results().empty();
+
+  p << AssumingOp::getOperationName() << " " << op.witness();
+  if (yieldsResults) {
+    p << " -> (" << op.getResultTypes() << ")";
+  }
+  p.printRegion(op.doRegion(),
+                /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/yieldsResults);
+  p.printOptionalAttrDict(op.getAttrs());
+}
+
+//===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 BroadcastOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,
                               SmallVectorImpl<Type> &inferredReturnTypes) {

mlir/test/Dialect/Shape/ops.mlir

 }
 
 func @test_constraints() {
   %0 = shape.const_shape []
   %1 = shape.const_shape [1, 2, 3]
   %w0 = "shape.cstr_broadcastable"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.witness
   %w1 = "shape.cstr_eq"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.witness
   %w3 = "shape.assuming_all"(%w0, %w1) : (!shape.witness, !shape.witness) -> !shape.witness
-  "shape.assuming"(%w3) ( {
+  shape.assuming %w3 -> !shape.shape {
     %2 = "shape.any"(%0, %1)  : (!shape.shape, !shape.shape) -> !shape.shape
     "shape.assuming_yield"(%2) : (!shape.shape) -> ()
-  }) : (!shape.witness) -> !shape.shape
+  }
   return
 }