[mlir][Linalg] Make Fill operation work on tensors.

[mravishankar]

Summary

Depends on D95109

[mravishankar]

Adding test for tiling of fill op on tensor

[mravishankar]

Rebase

[rriddle]

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

179

[mravishankar]

Rebase and adress comments

[mravishankar]

Closed by commit 6e8ef3b76ab65960edd6ee99f387e75564d8d9db

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

def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", []> {
-  let arguments = (ins AnyStridedMemRef:$output,
+  let arguments = (ins AnyShaped:$output,
                    AnyTypeOf<[AnyFloat, AnySignlessInteger, AnyVector]>:$value);
+  let results = (outs Optional<AnyRankedTensor>:$result);
   let extraClassDeclaration = libraryCallName # [{
     ValueRange inputs() { return {}; }
     ValueRange outputs() { return getOperands().take_front(); }
 
     // Rank-polymorphic.
     //   filling_value -> O(ivs) with parallel iterators.
     ArrayAttr iterator_types() {
       unsigned nPar = getOutputShapedType(0).getRank();
       return Builder(getContext()).getStrArrayAttr(
         SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName()));
     }
 
     ArrayAttr indexing_maps() {
       MLIRContext *context = getContext();
       // filling_value -> O(ivs)
       return Builder(getContext()).getAffineMapArrayAttr({
           extractOrIdentityMap(llvm::None, getNumParallelLoops(), context)});
     }
 
     static std::function<void(Block &)> getRegionBuilder() {
       return nullptr;
     }
   }];
 
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` type(operands) (`->` type($result)^)?}];
+
+  let builders = [
+    OpBuilderDAG<(ins "Value":$output, "Value":$value)>
+  ];
+
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 /// A base class for pooling operation such as conv. The arguments must contain
 /// optional arguments `strides`, `dilations` and `padding` with following type:

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

 }
 
 ///////////////////// Operations defined with Tablegen /////////////////////////
 // For such operations that do not correspond to library calls (i.e. defined in
 // LinalgOps.td), we define an overloaded `print` function and a
 // parse`className` function.
 
 //===----------------------------------------------------------------------===//
+// FillOp
+//===----------------------------------------------------------------------===//
+
+void FillOp::build(OpBuilder &builder, OperationState &result, Value output,
+                   Value value) {
+  build(builder, result, output.getType().dyn_cast<RankedTensorType>(), output,
+        value);
+}
+
+//===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) {
   build(builder, result, resultTensorTypes, inputs, outputs,
effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());
 }
 
 static LogicalResult verify(FillOp op) {
   auto viewType = op.getOutputShapedType(0);
   auto fillType = op.value().getType();
   if (viewType.getElementType() != fillType)
     return op.emitOpError("expects fill type to match view elemental type");
+  if (!op.getNumResults() && !viewType.isa<MemRefType>()) {
+    return op.emitOpError(
+        "expected fill op with no result value to use memref type");
+  }
   return success();
 }
 
 void CopyOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), input(),
                        SideEffects::DefaultResource::get());

mlir/test/Dialect/Linalg/invalid.mlir

 {
   // expected-error @+1 {{expected dimension 1 of collapsed type to be static value of 20}}
   %0 = linalg.reshape %arg0
          [affine_map<(d0, d1, d2) -> (d0)>,
           affine_map<(d0, d1, d2) -> (d1, d2)>] :
        memref<?x4x5xf32> into memref<?x?xf32>
   return %0 : memref<?x?xf32>
 }
+
+// -----
+
+func @illegal_fill_tensor_no_return(%arg0 : index, %arg1 : index, %arg2 : f32)
+{
+  %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
+  // expected-error @+1 {{expected fill op with no result value to use memref type}}
+  linalg.fill(%0, %arg2) : tensor<?x?xf32>, f32
+}
+
+// -----
+
+func @illegal_fill_memref_with_return(%arg0 : memref<?x?xf32>, %arg1 : f32) -> memref<?x?xf32>
+{
+  // expected-error @+1 {{unexpected #results > #outputs}}
+  %0 = linalg.fill(%arg0, %arg1) : memref<?x?xf32>, f32 -> memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}
+
+// -----
+
+func @illegal_fill_memref_with_tensor_return
+  (%arg0 : memref<?x?xf32>, %arg1 : f32) -> tensor<?x?xf32>
+{
+  // expected-error @+1 {{unexpected #results > #outputs}}
+  %0 = linalg.fill(%arg0, %arg1) : memref<?x?xf32>, f32 -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+func @illegal_fill_tensor_with_memref_return
+  (%arg0 : tensor<?x?xf32>, %arg1 : f32) -> memref<?x?xf32>
+{
+  // expected-error @+1 {{expected type of operand #0 ('tensor<?x?xf32>') to match type of corresponding result ('memref<?x?xf32>')}}
+  %0 = linalg.fill(%arg0, %arg1) : tensor<?x?xf32>, f32 -> memref<?x?xf32>
+  return %0 : memref<?x?xf32>
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

%0 = linalg.reshape %arg0
     memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
   return %0 : memref<?x?x?xf32>
 }
 // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0)>
 // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1)>
 // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>
 //     CHECK: func @legal_collapsing_reshape_dynamic_memref
 //     CHECK:   linalg.reshape %{{.+}} [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+
+// -----
+
+func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
+  %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
+  %1 = linalg.fill(%0, %arg2) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}
+// CHECK: %{{.+}} = linalg.fill(%{{.+}}, %{{.+}}) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>