-Refactored the linalg op python wrapper creation with a helper function.

Author: shahidact

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

@@ -859,9 +859,10 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
       Variadic<AnyType>:$inputs,
       Variadic<AnyShaped>:$outputs,
       DefaultValuedOptionalAttr<
-       AffineMapArrayAttr,
-       "BatchMatmulOp::getDefaultIndexingMaps($_builder.getContext())"
-      >:$indexing_maps
+        AffineMapArrayAttr,
+        "BatchMatmulOp::getDefaultIndexingMaps($_builder.getContext())"
+      >:$indexing_maps,
+      DefaultValuedOptionalAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
     );
     let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
     let regions = (region AnyRegion:$region);
@@ -887,9 +888,10 @@ def BatchMatmulOp : LinalgStructuredBase_Op<"batch_matmul", !listconcat([AttrSiz
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
-            CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+            "Attribute":$cast, CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
       [{
         $_state.addOperands(operands);
+        $_state.addAttribute("cast", cast);
         $_state.addAttributes(attributes);
         $_state.addTypes(resultTensorTypes);
         (void)$_state.addRegion(),

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -3951,11 +3951,18 @@ void BatchMatmulOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
   RegionBuilderHelper helper(b, block);
   SmallVector<Value> yields;
 
+  TypeFn castVal = TypeFn::cast_signed;
+  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
+    return attr.getName() == "cast";
+  });
+  if (castIter != attrs.end()) {
+    if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))
+      castVal = attr.getValue();
+  }
+
   auto toType = block.getArgument(2).getType();
-  Value castValA =
-      helper.buildTypeFn(TypeFn::cast_signed, toType, block.getArgument(0));
-  Value castValB =
-      helper.buildTypeFn(TypeFn::cast_signed, toType, block.getArgument(1));
+  Value castValA = helper.buildTypeFn(castVal, toType, block.getArgument(0));
+  Value castValB = helper.buildTypeFn(castVal, toType, block.getArgument(1));
   Value mulVal = helper.buildBinaryFn(BinaryFn::mul, castValA, castValB);
   Value addVal =
       helper.buildBinaryFn(BinaryFn::add, block.getArgument(2), mulVal);

mlir/python/mlir/dialects/linalg/__init__.py

@@ -149,7 +149,8 @@ def __init__(
 generic = region_op(GenericOp_, terminator=YieldOp)
 
 
-def matmul(
+def create_op(
+    op_type,
     *ins: Union[Operation, OpView, Value],
     outs: Sequence[Union[Operation, OpView, Value]],
     indexing_maps: Optional[Sequence[AffineMapAttr]] = None,
@@ -161,7 +162,7 @@ def matmul(
     init = _get_op_result_or_value(outs[0])
     result_types = [init.type] if isinstance(init.type, RankedTensorType) else []
 
-    op = MatmulOp(
+    op = op_type(
         result_tensors=result_types,
         inputs=ins,
         outputs=[init],
@@ -172,45 +173,32 @@ def matmul(
     return op
 
 
-def contract(
+def matmul(
     *ins: Union[Operation, OpView, Value],
     outs: Sequence[Union[Operation, OpView, Value]],
-    indexing_maps: Sequence[AffineMapAttr],
+    indexing_maps: Optional[Sequence[AffineMapAttr]] = None,
     cast: Optional[Union[TypeFn, Attribute]] = None,
 ):
-    ins = [_get_op_result_or_value(input) for input in ins]
-    if len(outs) > 1:
-        raise ValueError(f"{outs=} must have length 1.")
-    init = _get_op_result_or_value(outs[0])
-    result_types = [init.type] if isinstance(init.type, RankedTensorType) else []
-
-    op = ContractOp(
-        result_tensors=result_types,
-        inputs=ins,
-        outputs=[init],
-        indexing_maps=indexing_maps,
-        cast=cast,
-    )
-    fill_builtin_region(op.operation)
-    return op
+    return create_op(MatmulOp, *ins, outs=outs, indexing_maps=indexing_maps, cast=cast)
 
 
 def batch_matmul(
     *ins: Union[Operation, OpView, Value],
     outs: Sequence[Union[Operation, OpView, Value]],
     indexing_maps: Optional[Sequence[AffineMapAttr]] = None,
+    cast: Optional[Union[TypeFn, Attribute]] = None,
 ):
-    ins = [_get_op_result_or_value(input) for input in ins]
-    if len(outs) > 1:
-        raise ValueError(f"{outs=} must have length 1.")
-    init = _get_op_result_or_value(outs[0])
-    result_types = [init.type] if isinstance(init.type, RankedTensorType) else []
+    return create_op(
+        BatchMatmulOp, *ins, outs=outs, indexing_maps=indexing_maps, cast=cast
+    )
 
-    op = BatchMatmulOp(
-        result_tensors=result_types,
-        inputs=ins,
-        outputs=[init],
-        indexing_maps=indexing_maps,
+
+def contract(
+    *ins: Union[Operation, OpView, Value],
+    outs: Sequence[Union[Operation, OpView, Value]],
+    indexing_maps: Sequence[AffineMapAttr],
+    cast: Optional[Union[TypeFn, Attribute]] = None,
+):
+    return create_op(
+        ContractOp, *ins, outs=outs, indexing_maps=indexing_maps, cast=cast
     )
-    fill_builtin_region(op.operation)
-    return op