custom printer/parser to keep affine_maps inlined with the bind shape…

… ops

include/torch-mlir/Dialect/Torch/IR/TorchOps.td

@@ -1396,9 +1396,7 @@ def Torch_BindSymbolicShapeOp : Torch_Op<"bind_symbolic_shape", []> {
     Builtin_AffineMapAttr:$shape_expressions
   );
   let results = (outs);
-  let assemblyFormat = [{
-    $operand `,` `[` $shape_symbols `]` `,` $shape_expressions attr-dict `:` type($operand)
-  }];
+  let hasCustomAssemblyFormat = 1;
 }
 
 #endif // TORCH_OPS

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -5026,3 +5026,49 @@ LogicalResult InitializeGlobalSlotsOp::verify() {
     return emitOpError("expected number of operands to match number of slots");
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// BindSymbolicShapeOp
+//===----------------------------------------------------------------------===//
+
+//
+// torch.bind_symbolic_shape %6, [%0, %1, %2], #affine_map<()[s0, s1, s2] ->
+// (s0, 3, s1 * 2 + s2)> : !torch.vtensor<[?,3,?],f32>
+//
+
+ParseResult BindSymbolicShapeOp::parse(OpAsmParser &parser,
+                                       OperationState &result) {
+  OpAsmParser::UnresolvedOperand operand;
+  SmallVector<OpAsmParser::UnresolvedOperand> shapeSymbols;
+  AffineMapAttr shapeExpressions;
+  Type tensorType;
+
+  if (parser.parseOperand(operand) || parser.parseComma() ||
+      parser.parseLSquare() || parser.parseOperandList(shapeSymbols) ||
+      parser.parseRSquare() || parser.parseComma() ||
+      parser.parseAttribute(shapeExpressions, "shape_expressions",
+                            result.attributes) ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(tensorType)) {
+    return failure();
+  }
+
+  result.addTypes(tensorType);
+  if (parser.resolveOperand(operand, tensorType, result.operands) ||
+      parser.resolveOperands(shapeSymbols,
+                             parser.getBuilder().getIntegerType(64),
+                             result.operands)) {
+    return failure();
+  }
+
+  return success();
+}
+
+void BindSymbolicShapeOp::print(OpAsmPrinter &p) {
+  p << " " << getOperand() << ", [";
+  llvm::interleaveComma(getShapeSymbols(), p);
+  p << "], " << "#affine_map<" << getShapeExpressions().getValue() << ">";
+  p.printOptionalAttrDict((*this)->getAttrs(),
+                          /*elidedAttrs=*/{"shape_expressions"});
+  p << " : " << getOperand().getType();
+}