[mlir] Use assemblyFormat in AllocLikeOp.

[csigg]

Summary

Split operands into dynamicSizes and symbolOperands.

[bondhugula]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

Do you need to insert a space? This would make it inconsistent with affine.apply and perhaps other things as well. Is it necessary to get the assembly format working? If it's for readability, nearly all coding styles (for eg. MLIR, TensorFlow C++ styles) don't put a space in such cases between a bracket list and the parenthetical list - for eg. in lambda definitions. The square brackets were mainly used to denote a "capture" like effect.

[csigg]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

The OpFormatGen produces a space between two punctuations, unless the second one is a closing bracket or a comma.

I've tried to change that behavior to no space before any bracket, but that would require inserting explicit spaces in many other assemblyFormats. I've also tried to parse dynamicSizes and symbolOperands as a single custom group, but that inserts a space after alloc.

Does the 'no space between )[' apply almost everywhere so that we should handle that case differentely?

Maybe the best alternative is to add an empty literal that would set shouldEmitSpace = false. Basically the opposite of the space literal.

@rriddle, WDYT?

[csigg]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

I've made affine ops consistent by adding a space too. The code change itself is small, but I needed to update 80+ tests. Maybe it's easier to support a `` literal in the assemblyFormat: D91068.

[bondhugula]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

I think we could leave it the way it was before. If you are moving this to assembly format and it requires a space in the syntax, let's go with it. It's a bit counterintuitive for those writing MLIR test cases to know that the space there is making a difference.

[csigg]

Also add space to affine_map and affine.apply etc.

[bondhugula]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

I'm sorry if it wasn't clear. I'm not in favor of adding a space for all the affine ops - it's unnecessary and counterintuitive (also inconsistent with the way affine maps defs are printed - space shouldn't be making a difference). Since you are moving just the alloc ops to the assembly format, it's fine if only the alloc op use a space in the syntax. The alloc ops and affine ops need not really be consistent since it's at the expense of other downsides.

[csigg]

mlir/test/IR/memory-ops.mlir

25 ↗

(On Diff #302219)

Is it inconsistent with affine maps? I think I updated that too.

I don't immediately see why an extra space is counterintuitive, but I definitely see the unnecessary churn on all the lit tests, just because assemblyFormat happens to insert a space.

I will see how D91068 is going, and roll this one back to just the std.alloc change, either with or without space.

[csigg]

Don't introduce extra space.

[stella.stamenova]

onnx-mlir (https://github.com/onnx/onnx-mlir) is currently using parseDimAndSymbolList. Is the intention for this function to not be available outside of mlir any more? I notice it no longer has any uses inside mlir either.

[csigg]

In D90589#2391801, @stella.stamenova wrote:

onnx-mlir (https://github.com/onnx/onnx-mlir) is currently using parseDimAndSymbolList. Is the intention for this function to not be available outside of mlir any more? I notice it no longer has any uses inside mlir either.

It's now private to AffineOps.cpp. Maybe the easiest would be to export it again from there. Would that work for onnx?

[stella.stamenova]

In D90589#2393110, @csigg wrote:

In D90589#2391801, @stella.stamenova wrote:

onnx-mlir (https://github.com/onnx/onnx-mlir) is currently using parseDimAndSymbolList. Is the intention for this function to not be available outside of mlir any more? I notice it no longer has any uses inside mlir either.

It's now private to AffineOps.cpp. Maybe the easiest would be to export it again from there. Would that work for onnx?

Yes, that should. Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

 /// `originalShape` to obtain the `reducedShape`. This mask is useful to track
 /// which dimensions must be kept when e.g. compute MemRef strides under
 /// rank-reducing operations. Return None if reducedShape cannot be obtained
 /// by dropping only `1` entries in `originalShape`.
 llvm::Optional<SmallVector<bool, 4>>
 computeRankReductionMask(ArrayRef<int64_t> originalShape,
                          ArrayRef<int64_t> reducedShape);
 
-/// Prints dimension and symbol list.
-void printDimAndSymbolList(Operation::operand_iterator begin,
-                           Operation::operand_iterator end, unsigned numDims,
-                           OpAsmPrinter &p);
-
-/// Parses dimension and symbol list and returns true if parsing failed.
-ParseResult parseDimAndSymbolList(OpAsmParser &parser,
-                                  SmallVectorImpl<Value> &operands,
-                                  unsigned &numDims);
-
 /// Determines whether MemRefCastOp casts to a more dynamic version of the
 /// source memref. This is useful to to fold a memref_cast into a consuming op
 /// and implement canonicalization patterns for ops in different dialects that
 /// may consume the results of memref_cast operations. Such foldable memref_cast
 /// operations are typically inserted as `view` and `subview` ops and are
 /// canonicalized, to preserve the type compatibility of their uses.
 ///
 /// Returns true when all conditions are met:

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 
 // Base class for memref allocating ops: alloca and alloc.
 //
 //   %0 = alloclike(%m)[%s] : memref<8x?xf32, (d0, d1)[s0] -> ((d0 + s0), d1)>
 //
 class AllocLikeOp<string mnemonic,
                   Resource resource,
                   list<OpTrait> traits = []> :
-    Std_Op<mnemonic, !listconcat([MemoryEffects<[MemAlloc<resource>]>], traits)> {
-
-  let arguments = (ins Variadic<Index>:$value,
+    Std_Op<mnemonic,
+    !listconcat([
+      MemoryEffects<[MemAlloc<resource>]>,
+      AttrSizedOperandSegments
+    ], traits)> {
+
+  let arguments = (ins Variadic<Index>:$dynamicSizes,
+                   // The symbolic operands (the ones in square brackets) bind
+                   // to the symbols of the memref's layout map.
+                   Variadic<Index>:$symbolOperands,
                    Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment);
-  let results = (outs Res<AnyMemRef, "", [MemAlloc<resource>]>);
+  let results = (outs Res<AnyMemRef, "", [MemAlloc<resource>]>:$memref);
 
   let builders = [
-    OpBuilderDAG<(ins "MemRefType":$memrefType), [{
-      $_state.types.push_back(memrefType);
+    OpBuilderDAG<(ins "MemRefType":$memrefType,
+                  CArg<"IntegerAttr", "IntegerAttr()">:$alignment), [{
+      return build($_builder, $_state, memrefType, {}, alignment);
     }]>,
-    OpBuilderDAG<(ins "MemRefType":$memrefType, "ValueRange":$operands,
+    OpBuilderDAG<(ins "MemRefType":$memrefType, "ValueRange":$dynamicSizes,
                   CArg<"IntegerAttr", "IntegerAttr()">:$alignment), [{
-      $_state.addOperands(operands);
+      return build($_builder, $_state, memrefType, dynamicSizes, {}, alignment);
+    }]>,
+    OpBuilderDAG<(ins "MemRefType":$memrefType, "ValueRange":$dynamicSizes,
+                  "ValueRange":$symbolOperands,
+                  CArg<"IntegerAttr", "{}">:$alignment), [{
       $_state.types.push_back(memrefType);
+      $_state.addOperands(dynamicSizes);
+      $_state.addOperands(symbolOperands);
+      $_state.addAttribute(getOperandSegmentSizeAttr(),
+          $_builder.getI32VectorAttr({
+              static_cast<int32_t>(dynamicSizes.size()),
+              static_cast<int32_t>(symbolOperands.size())}));
       if (alignment)
         $_state.addAttribute(getAlignmentAttrName(), alignment);
     }]>];
 
   let extraClassDeclaration = [{
     static StringRef getAlignmentAttrName() { return "alignment"; }
 
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
 
-    /// Returns the number of symbolic operands (the ones in square brackets),
-    /// which bind to the symbols of the memref's layout map.
-    unsigned getNumSymbolicOperands() {
-      return getNumOperands() - getType().getNumDynamicDims();
-    }
-
-    /// Returns the symbolic operands (the ones in square brackets), which bind
-    /// to the symbols of the memref's layout map.
-    operand_range getSymbolicOperands() {
-      return {operand_begin() + getType().getNumDynamicDims(), operand_end()};
-    }
-
     /// Returns the dynamic sizes for this alloc operation if specified.
-    operand_range getDynamicSizes() { return getOperands(); }
+    operand_range getDynamicSizes() { return dynamicSizes(); }
   }];
 
-  let parser = [{ return ::parseAllocLikeOp(parser, result); }];
+  let assemblyFormat = [{
+    `(`$dynamicSizes`)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
+  }];
 
   let hasCanonicalizer = 1;
 }
 
 // Base class for ops with static/dynamic offset, sizes and strides
 // attributes/arguments.
 class BaseOpWithOffsetSizesAndStrides<string mnemonic, list<OpTrait> traits = []> :
     Std_Op<mnemonic,

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 // Returns true if 'value' is a valid index to an affine operation (e.g.
 // affine.load, affine.store, affine.dma_start, affine.dma_wait) where
 // `region` provides the polyhedral symbol scope. Returns false otherwise.
 static bool isValidAffineIndexOperand(Value value, Region *region) {
   return isValidDim(value, region) || isValidSymbol(value, region);
 }
 
+/// Prints dimension and symbol list.
+static void printDimAndSymbolList(Operation::operand_iterator begin,
+                                  Operation::operand_iterator end,
+                                  unsigned numDims, OpAsmPrinter &printer) {
+  OperandRange operands(begin, end);
+  printer << '(' << operands.take_front(numDims) << ')';
+  if (operands.size() > numDims)
+    printer << '[' << operands.drop_front(numDims) << ']';
+}
+
+/// Parses dimension and symbol list and returns true if parsing failed.
+static ParseResult parseDimAndSymbolList(OpAsmParser &parser,
+                                         SmallVectorImpl<Value> &operands,
+                                         unsigned &numDims) {
+  SmallVector<OpAsmParser::OperandType, 8> opInfos;
+  if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren))
+    return failure();
+  // Store number of dimensions for validation by caller.
+  numDims = opInfos.size();
+
+  // Parse the optional symbol operands.
+  auto indexTy = parser.getBuilder().getIndexType();
+  return failure(parser.parseOperandList(
+                     opInfos, OpAsmParser::Delimiter::OptionalSquare) ||
+                 parser.resolveOperands(opInfos, indexTy, operands));
+}
+
 /// Utility function to verify that a set of operands are valid dimension and
 /// symbol identifiers. The operands should be laid out such that the dimension
 /// operands are before the symbol operands. This function returns failure if
 /// there was an invalid operand. An operation is provided to emit any necessary
 /// errors.
 template <typename OpTy>
 static LogicalResult
 verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands,

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *StandardOpsDialect::materializeConstant(OpBuilder &builder,
                                                    Attribute value, Type type,
                                                    Location loc) {
   return builder.create<ConstantOp>(loc, type, value);
 }
 
-void mlir::printDimAndSymbolList(Operation::operand_iterator begin,
-                                 Operation::operand_iterator end,
-                                 unsigned numDims, OpAsmPrinter &p) {
-  Operation::operand_range operands(begin, end);
-  p << '(' << operands.take_front(numDims) << ')';
-  if (operands.size() != numDims)
-    p << '[' << operands.drop_front(numDims) << ']';
-}
-
-// Parses dimension and symbol list, and sets 'numDims' to the number of
-// dimension operands parsed.
-// Returns 'false' on success and 'true' on error.
-ParseResult mlir::parseDimAndSymbolList(OpAsmParser &parser,
-                                        SmallVectorImpl<Value> &operands,
-                                        unsigned &numDims) {
-  SmallVector<OpAsmParser::OperandType, 8> opInfos;
-  if (parser.parseOperandList(opInfos, OpAsmParser::Delimiter::Paren))
-    return failure();
-  // Store number of dimensions for validation by caller.
-  numDims = opInfos.size();
-
-  // Parse the optional symbol operands.
-  auto indexTy = parser.getBuilder().getIndexType();
-  if (parser.parseOperandList(opInfos,
-                              OpAsmParser::Delimiter::OptionalSquare) ||
-      parser.resolveOperands(opInfos, indexTy, operands))
-    return failure();
-  return success();
-}
-
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses m_Constant
 /// and checks the operation for an index type.
 static detail::op_matcher<ConstantIndexOp> m_ConstantIndex() {
   return detail::op_matcher<ConstantIndexOp>();
 }
 
 //===----------------------------------------------------------------------===//
static LogicalResult verifyOpWithOffsetSizesAndStrides(OpType op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp / AllocaOp
 //===----------------------------------------------------------------------===//
 
 template <typename AllocLikeOp>
-static void printAllocLikeOp(OpAsmPrinter &p, AllocLikeOp op, StringRef name) {
-  static_assert(llvm::is_one_of<AllocLikeOp, AllocOp, AllocaOp>::value,
-                "applies to only alloc or alloca");
-  p << name;
-
-  // Print dynamic dimension operands.
-  MemRefType type = op.getType();
-  printDimAndSymbolList(op.operand_begin(), op.operand_end(),
-                        type.getNumDynamicDims(), p);
-  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"map"});
-  p << " : " << type;
-}
-
-static void print(OpAsmPrinter &p, AllocOp op) {
-  printAllocLikeOp(p, op, "alloc");
-}
-
-static void print(OpAsmPrinter &p, AllocaOp op) {
-  printAllocLikeOp(p, op, "alloca");
-}
-
-static ParseResult parseAllocLikeOp(OpAsmParser &parser,
-                                    OperationState &result) {
-  MemRefType type;
-
-  // Parse the dimension operands and optional symbol operands, followed by a
-  // memref type.
-  unsigned numDimOperands;
-  if (parseDimAndSymbolList(parser, result.operands, numDimOperands) ||
-      parser.parseOptionalAttrDict(result.attributes) ||
-      parser.parseColonType(type))
-    return failure();
-
-  // Check numDynamicDims against number of question marks in memref type.
-  // Note: this check remains here (instead of in verify()), because the
-  // partition between dim operands and symbol operands is lost after parsing.
-  // Verification still checks that the total number of operands matches
-  // the number of symbols in the affine map, plus the number of dynamic
-  // dimensions in the memref.
-  if (numDimOperands != type.getNumDynamicDims())
-    return parser.emitError(parser.getNameLoc())
-           << "dimension operand count does not equal memref dynamic dimension "
-              "count";
-  result.types.push_back(type);
-  return success();
-}
-
-template <typename AllocLikeOp>
-static LogicalResult verify(AllocLikeOp op) {
+static LogicalResult verifyAllocLikeOp(AllocLikeOp op) {
   static_assert(llvm::is_one_of<AllocLikeOp, AllocOp, AllocaOp>::value,
                 "applies to only alloc or alloca");
   auto memRefType = op.getResult().getType().template dyn_cast<MemRefType>();
   if (!memRefType)
     return op.emitOpError("result must be a memref");
 
-  unsigned numSymbols = 0;
-  if (!memRefType.getAffineMaps().empty()) {
-    // Store number of symbols used in affine map (used in subsequent check).
-    AffineMap affineMap = memRefType.getAffineMaps()[0];
-    numSymbols = affineMap.getNumSymbols();
-  }
+  if (static_cast<int64_t>(op.dynamicSizes().size()) !=
+      memRefType.getNumDynamicDims())
+    return op.emitOpError("dimension operand count does not equal memref "
+                          "dynamic dimension count");
 
-  // Check that the total number of operands matches the number of symbols in
-  // the affine map, plus the number of dynamic dimensions specified in the
-  // memref type.
-  unsigned numDynamicDims = memRefType.getNumDynamicDims();
-  if (op.getNumOperands() != numDynamicDims + numSymbols)
+  unsigned numSymbols = 0;
+  if (!memRefType.getAffineMaps().empty())
+    numSymbols = memRefType.getAffineMaps().front().getNumSymbols();
+  if (op.symbolOperands().size() != numSymbols)
     return op.emitOpError(
-        "operand count does not equal dimension plus symbol operand count");
+        "symbol operand count does not equal memref symbol count");
 
-  // Verify that all operands are of type Index.
-  for (auto operandType : op.getOperandTypes())
-    if (!operandType.isIndex())
-      return op.emitOpError("requires operands to be of type Index");
-
-  if (std::is_same<AllocLikeOp, AllocOp>::value)
-    return success();
+}
 
+static LogicalResult verify(AllocOp op) { return verifyAllocLikeOp(op); }
+
+static LogicalResult verify(AllocaOp op) {
   // An alloca op needs to have an ancestor with an allocation scope trait.
-  if (!op.template getParentWithTrait<OpTrait::AutomaticAllocationScope>())
+  if (!op.getParentWithTrait<OpTrait::AutomaticAllocationScope>())
     return op.emitOpError(
         "requires an ancestor op with AutomaticAllocationScope trait");
 
-  return success();
+  return verifyAllocLikeOp(op);
 }
 
 namespace {
 /// Fold constant dimensions into an alloc like operation.
 template <typename AllocLikeOp>
 struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
   using OpRewritePattern<AllocLikeOp>::OpRewritePattern;
 

mlir/lib/Transforms/Utils/Utils.cpp

 // TODO: Currently works for static memrefs with a single layout map.
 LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
   MemRefType memrefType = allocOp.getType();
   OpBuilder b(allocOp);
 
   // Fetch a new memref type after normalizing the old memref to have an
   // identity map layout.
   MemRefType newMemRefType =
-      normalizeMemRefType(memrefType, b, allocOp.getNumSymbolicOperands());
+      normalizeMemRefType(memrefType, b, allocOp.symbolOperands().size());
   if (newMemRefType == memrefType)
     // Either memrefType already had an identity map or the map couldn't be
     // transformed to an identity map.
     return failure();
 
   Value oldMemRef = allocOp.getResult();
 
-  SmallVector<Value, 4> symbolOperands(allocOp.getSymbolicOperands());
+  SmallVector<Value, 4> symbolOperands(allocOp.symbolOperands());
   AllocOp newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType,
-                                       llvm::None, allocOp.alignmentAttr());
+                                       allocOp.alignmentAttr());
   AffineMap layoutMap = memrefType.getAffineMaps().front();
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
                                       /*extraOperands=*/{},
                                       /*symbolOperands=*/symbolOperands,
                                       /*domInstFilter=*/nullptr,

mlir/test/IR/invalid-ops.mlir

 }
 
 // -----
 
 func @bad_alloc_wrong_dynamic_dim_count() {
 ^bb0:
   %0 = constant 7 : index
   // Test alloc with wrong number of dynamic dimensions.
-  %1 = alloc(%0)[%1] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1> // expected-error {{op 'std.alloc' dimension operand count does not equal memref dynamic dimension count}}
+  // expected-error@+1 {{dimension operand count does not equal memref dynamic dimension count}}
+  %1 = alloc(%0)[%0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
   return
 }
 
 // -----
 
 func @bad_alloc_wrong_symbol_count() {
 ^bb0:
   %0 = constant 7 : index
   // Test alloc with wrong number of symbols
-  %1 = alloc(%0) : memref<2x?xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1> // expected-error {{operand count does not equal dimension plus symbol operand count}}
+  // expected-error@+1 {{symbol operand count does not equal memref symbol count}}
+  %1 = alloc(%0) : memref<2x?xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
   return
 }
 
 // -----
 
 func @test_store_zero_results() {
 ^bb0:
   %0 = alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>