[mlir][linalg] Add a FillOpInterface.

[gysit]

Summary

Add a FillOpInterface similar to the contraction and convolution op interfaces. The FillOpInterface is a preparation step to replace linalg.fill by its OpDSL version linalg.fill_tensor. The interface implements the value(), output(), and result() methods that by default are not available on linalg.fill_tensor.

[gysit]

Rebase.

[gysit]

Regenerate yaml file.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h

 namespace detail {
 
 /// Verify that `op` conforms to ContractionOpInterface.
 LogicalResult verifyContractionInterface(Operation *op);
 
 /// Verify that `op` conforms to the ConvolutionOpInterface.
 LogicalResult verifyConvolutionInterface(Operation *op);
 
+/// Verify that `op` conforms to the FillOpInterface.
+LogicalResult verifyFillInterface(Operation *op);
+
 /// Verify that `op` conforms to the invariants of StructuredOpInterface
 LogicalResult verifyStructuredOpInterface(Operation *op);
 
 } // namespace detail
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc"
 
 /// Include the generated interface declarations.
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h.inc"
 
 #endif // MLIR_DIALECT_LINALG_IR_LINALGINTERFACES_H_

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

InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getOperation()->getOperand(1);
       }]
     >,
   ];
 }
 
+def LinalgFillOpInterface : OpInterface<"FillOpInterface"> {
+  let description = [{
+    A fill operation is defined in general terms:
+    1. Has a scalar `value` operand.
+    2. Has one `output` operand.
+  }];
+  let cppNamespace = "::mlir::linalg";
+  let verify = [{ return detail::verifyFillInterface($_op); }];
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Return the fill value.",
+      /*retTy=*/"Value",
+      /*methodName=*/"value",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getOperation()->getOperand(0);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the output operand.",
+      /*retTy=*/"Value",
+      /*methodName=*/"output",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getOperation()->getOperand(1);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the result.",
+      /*retTy=*/"Value",
+      /*methodName=*/"result",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        if ($_op.getOperation()->getResults().empty())
+          return nullptr;
+        return $_op.getOperation()->getResults().front();
+      }]
+    >,
+  ];
+}
+
 // The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
 def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
   let cppNamespace = "::mlir::linalg";
   let methods = [
     //===------------------------------------------------------------------===//
     // Loop types handling.
     //===------------------------------------------------------------------===//
     InterfaceMethod<

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

metadata: !LinalgOpMetadata
   name: fill_tensor
   cpp_class_name: FillTensorOp
   doc: |-
     Fills the output tensor with the given value.
 
     Works for arbitrary ranked output tensors since the operation performs scalar
     accesses only and is thus rank polymorphic. Numeric casting is performed on
     the value operand, promoting it to the same data type as the output.
+  implements:
+  - LinalgFillOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: value
     kind: scalar
     type_var: T1
   - !LinalgOperandDefConfig
     name: O

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

return op->emitError(
         "expected all iterators used to access outputs to be parallel");
   }
   if (res == MatchConvolutionResult::NonOutputDimNotReduction) {
     return op->emitError(
         "expected all iterators not used to access outputs to be reduction");
   }
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// FillOpInterface implementation
+//===----------------------------------------------------------------------===//
+
+enum class MatchFillResult {
+  Success = 0,
+  NotLinalgOp,
+  WrongNumOperands,
+  NotScalarInput
+};
+
+static MatchFillResult isFillInterfaceImpl(Operation *op) {
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+  if (!linalgOp)
+    return MatchFillResult::NotLinalgOp;
+  if (linalgOp.getNumInputs() != 1 || linalgOp.getNumOutputs() != 1)
+    return MatchFillResult::WrongNumOperands;
+
+  OpOperand *value = linalgOp.getInputOperand(0);
+  if (!linalgOp.isScalar(value))
+    return MatchFillResult::NotScalarInput;
+
+  return MatchFillResult::Success;
+}
+
+LogicalResult mlir::linalg::detail::verifyFillInterface(Operation *op) {
+  auto res = isFillInterfaceImpl(op);
+  if (res == MatchFillResult::NotLinalgOp)
+    return op->emitError("expected a LinalgOp");
+  if (res == MatchFillResult::WrongNumOperands)
+    return op->emitError("expected op with 1 input and 1 output");
+  if (res == MatchFillResult::NotScalarInput)
+    return op->emitError("expected op with scalar input");
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // StructuredOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 OpOperandVector::operator SmallVector<Value>() {
   SmallVector<Value> result;
   result.reserve(this->size());
   llvm::transform(*this, std::back_inserter(result),

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

class OpInterfaceDef:
   """An interface that an op implements."""
 
   def __init__(self, cpp_name: str):
     self.cpp_name = cpp_name
 
 
 ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
 ConvolutionOpInterface = OpInterfaceDef("LinalgConvolutionOpInterface")
+FillOpInterface = OpInterfaceDef("LinalgFillOpInterface")
 
 
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
   def __init__(self, name: str, cpp_class_name: Optional[str],
                doc: Optional[str]):

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

 @linalg_structured_op
 def fill_tensor(value=ScalarDef(T1), O=TensorDef(U, output=True)):
   """Fills the output tensor with the given value.
 
   Works for arbitrary ranked output tensors since the operation performs scalar
   accesses only and is thus rank polymorphic. Numeric casting is performed on
   the value operand, promoting it to the same data type as the output.
   """
+  implements(FillOpInterface)
   O[None] = TypeFn.cast_signed(U, value)
 
 
 @linalg_structured_op
 def fill_rng_2d(
     min=ScalarDef(F64),
     max=ScalarDef(F64),
     seed=ScalarDef(I32),

mlir/test/Dialect/Linalg/fill-interface-invalid.mlir

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+func @test_fill_op_not_linalg_op(%arg0 : f32, %arg1 : tensor<?xf32>)
+     -> tensor<?xf32> {
+  // expected-error @+1 {{expected a LinalgOp}}
+  %0 = "test.fill_op_not_linalg_op"(%arg0, %arg1)
+      : (f32, tensor<?xf32>) -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+#map0 = affine_map<(d0) -> ()>
+#map1 = affine_map<(d0) -> (d0)>
+func @test_fill_op_wrong_num_operands(%arg0 : f32, %arg1 : tensor<?xf32>)
+     -> tensor<?xf32> {
+  // expected-error @+1 {{expected op with 1 input and 1 output}}
+  %0 = test.linalg_fill_op {
+      indexing_maps = [#map0, #map0, #map1],
+      iterator_types = ["parallel"]}
+      ins(%arg0, %arg0 : f32, f32) outs(%arg1 : tensor<?xf32>) {
+      ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+         linalg.yield  %arg2 : f32
+      } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+#map1 = affine_map<(d0) -> (d0)>
+func @test_fill_op_non_scalar_input(%arg0 : tensor<?xf32>,
+    %arg1 : tensor<?xf32>) -> tensor<?xf32> {
+  // expected-error @+1 {{expected op with scalar input}}
+  %0 = test.linalg_fill_op {
+      indexing_maps = [#map1, #map1],
+      iterator_types = ["parallel"]}
+      ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<?xf32>) {
+      ^bb0(%arg2 : f32, %arg3 : f32):
+         linalg.yield  %arg2 : f32
+      } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}

mlir/test/lib/Dialect/Test/TestOps.td

let extraClassDeclaration = [{
 
     // To conform with interface requirement on operand naming.
     mlir::ValueRange inputs() { return getInputs(); }
     mlir::ValueRange outputs() { return getOutputs(); }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// Test LinalgFillOpInterface.
+//===----------------------------------------------------------------------===//
+
+def TestLinalgFillOpNotLinalgOp : TEST_Op<"fill_op_not_linalg_op", [
+    LinalgFillOpInterface]> {
+  let arguments = (ins
+    AnyType:$value, AnyType:$output);
+  let results = (outs AnyRankedTensor:$result);
+}
+
+def TestLinalgFillOp :
+  TEST_Op<"linalg_fill_op", [AttrSizedOperandSegments, SingleBlock,
+      LinalgStructuredInterface, LinalgFillOpInterface]> {
+
+  let arguments = (ins Variadic<AnyType>:$inputs,
+    Variadic<AnyType>:$outputs);
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region AnyRegion:$region);
+
+  let assemblyFormat = [{
+    attr-dict (`ins` `(` $inputs^ `:` type($inputs) `)`)?
+    `outs` `(` $outputs `:` type($outputs) `)`
+    $region (`->` type($results)^)?
+  }];
+
+  let extraClassDeclaration = [{
+    bool hasIndexSemantics() { return false; }
+
+    static void regionBuilder(mlir::ImplicitLocOpBuilder &b, mlir::Block &block,
+                              mlir::ArrayRef<mlir::NamedAttribute> attrs) {
+      b.create<mlir::linalg::YieldOp>(block.getArguments().back());
+    }
+
+    static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
+                              mlir::ArrayRef<mlir::NamedAttribute>)>
+    getRegionBuilder() {
+      return &regionBuilder;
+    }
+
+    mlir::ArrayAttr iterator_types() {
+      return getOperation()->getAttrOfType<mlir::ArrayAttr>("iterator_types");
+    }
+
+    mlir::ArrayAttr indexing_maps() {
+      return getOperation()->getAttrOfType<mlir::ArrayAttr>("indexing_maps");
+    }
+
+    std::string getLibraryCallName() {
+      return "";
+    }
+
+    // To conform with interface requirement on operand naming.
+    mlir::ValueRange inputs() { return getInputs(); }
+    mlir::ValueRange outputs() { return getOutputs(); }
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // Test Ops with Default-Valued String Attributes
 //===----------------------------------------------------------------------===//
 
 def TestDefaultStrAttrNoValueOp : TEST_Op<"no_str_value"> {
   let arguments = (ins DefaultValuedAttr<StrAttr, "">:$value);
   let assemblyFormat = "attr-dict";
 }