[mlir][shape] refine shape.func and shape.with_shape

[yaochengji]

Summary

• shape.with_shape supports ExtentTensorType
• add helper to create shape.func

[jpienaar]

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

1293

getSymNameAttr ? They do match (on purpose) but sym_name is the attribute name ODS defined and what the autogenerated builders use.

1295

Similiarly in the generated code we have

odsState.addAttribute(getFunctionTypeAttrName(odsState.name), function_type);

rather than querying the generic. (the difference is the latter form uses a cached Attribute rather than looking up the stringref and uniqueing it).

[yaochengji]

Thanks for pointing this out, use getSymNameAttrName and getFunctionTypeAttrName instead

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

let description = [{
     This op need not be a refinement of the shape. In non-error cases the input
     ValueShape's value and shape are conformant and so too for the output, but
     the result may be less specified than `operand`'s shape as `shape` is
     merely used to construct the new ValueShape. If join behavior is desired
     then a join op should be used.
   }];
 
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$operand,
-                       Shape_ShapeType:$shape);
+                       Shape_ShapeOrExtentTensorType:$shape);
   let results = (outs Shape_ValueShapeType:$result);
 
   let assemblyFormat = "operands attr-dict `:` type($operand) `,` type($shape)";
 }
 
 def Shape_YieldOp : Shape_Op<"yield",
     [HasParent<"ReduceOp, FunctionLibraryOp">,
      NoSideEffect,
let description = [{
     function.
   }];
 
   let arguments = (ins SymbolNameAttr:$sym_name,
                        TypeAttrOf<FunctionType>:$function_type,
                        OptionalAttr<StrAttr>:$sym_visibility);
   let regions = (region AnyRegion:$body);
 
+  let builders = [OpBuilder<(ins
+    "StringRef":$name, "FunctionType":$type,
+    CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
+    CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)
+  >];
+
   let extraClassDeclaration = [{
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         ArrayRef<NamedAttribute> attrs = {});
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         Operation::dialect_attr_range attrs);
+    static FuncOp create(Location location, StringRef name, FunctionType type,
+                         ArrayRef<NamedAttribute> attrs,
+                         ArrayRef<DictionaryAttr> argAttrs);
     //===------------------------------------------------------------------===//
     // CallableOpInterface
     //===------------------------------------------------------------------===//
 
     /// Returns the region on the current operation that is callable. This may
     /// return null in the case of an external callable object, e.g. an external
     /// function.
     ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }

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

void FunctionLibraryOp::print(OpAsmPrinter &p) {
   p << " mapping ";
   p.printAttributeWithoutType(getMappingAttr());
 }
 
 //===----------------------------------------------------------------------===//
 // FuncOp
 //===----------------------------------------------------------------------===//
 
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      ArrayRef<NamedAttribute> attrs) {
+  OpBuilder builder(location->getContext());
+  OperationState state(location, getOperationName());
+  FuncOp::build(builder, state, name, type, attrs);
+  return cast<FuncOp>(Operation::create(state));
+}
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      Operation::dialect_attr_range attrs) {
+  SmallVector<NamedAttribute, 8> attrRef(attrs);
+  return create(location, name, type, llvm::makeArrayRef(attrRef));
+}
+FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
+                      ArrayRef<NamedAttribute> attrs,
+                      ArrayRef<DictionaryAttr> argAttrs) {
+  FuncOp func = create(location, name, type, attrs);
+  func.setAllArgAttrs(argAttrs);
+  return func;
+}
+
+void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
+                   FunctionType type, ArrayRef<NamedAttribute> attrs,
+                   ArrayRef<DictionaryAttr> argAttrs) {
+  state.addAttribute(FuncOp::getSymNameAttrName(state.name),
+                     builder.getStringAttr(name));
+  state.addAttribute(FuncOp::getFunctionTypeAttrName(state.name),
+                     TypeAttr::get(type));
+  state.attributes.append(attrs.begin(), attrs.end());
+  state.addRegion();
+
+  if (argAttrs.empty())
+    return;
+  assert(type.getNumInputs() == argAttrs.size());
+  function_interface_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                                /*resultAttrs=*/llvm::None);
+}
+
 ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {
   auto buildFuncType =
       [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
          function_interface_impl::VariadicFlag,
          std::string &) { return builder.getFunctionType(argTypes, results); };
 
   return function_interface_impl::parseFunctionOp(
       parser, result, /*allowVariadic=*/false, buildFuncType);

mlir/test/Dialect/Shape/ops.mlir

 }
 func.func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
   %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
   %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
   %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
   return %2 : !shape.shape
 }
 
+func.func @shape_with_shape_extent_tensor_type(%a : tensor<?x?x?xf32>, %b : !shape.value_shape) -> !shape.value_shape {
+  %0 = shape.shape_of %a : tensor<?x?x?xf32> -> tensor<3xindex>
+  %1 = shape.with_shape %b, %0 : !shape.value_shape, tensor<3xindex>
+  return %1 : !shape.value_shape
+}
+
 func.func @any_on_shape(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
     -> !shape.shape {
   %result = shape.any %a, %b, %c
       : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
   return %result : !shape.shape
 }
 
 func.func @any_on_mixed(%a : tensor<?xindex>,