Skip to content

Commit

Permalink
[mlir][Linalg] Update the semantics, verifier and test for Linalg wit…
Browse files Browse the repository at this point in the history 
…h tensors.

Summary:
This diff fixes issues with the semantics of linalg.generic on tensors that appeared when converting directly from HLO to linalg.generic.
The changes are self-contained within MLIR and can be captured and tested independently of XLA.

The linalg.generic and indexed_generic are updated to:

To allow progressive lowering from the value world (a.k.a tensor values) to
the buffer world (a.k.a memref values), a linalg.generic op accepts
mixing input and output ranked tensor values with input and output memrefs.

```
%1 = linalg.generic #trait_attribute %A, %B {other-attributes} :
  tensor<?x?xf32>,
  memref<?x?xf32, stride_specification>
  -> (tensor<?x?xf32>)
```

In this case, the number of outputs (args_out) must match the sum of (1) the
number of output buffer operands and (2) the number of tensor return values.
The semantics is that the linalg.indexed_generic op produces (i.e.
allocates and fills) its return values.

Tensor values must be legalized by a buffer allocation pass before most
transformations can be applied. Such legalization moves tensor return values
into output buffer operands and updates the region argument accordingly.

Transformations that create control-flow around linalg.indexed_generic
operations are not expected to mix with tensors because SSA values do not
escape naturally. Still, transformations and rewrites that take advantage of
tensor SSA values are expected to be useful and will be added in the near
future.

Subscribers: bmahjour, mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D72555
  • Loading branch information
Nicolas Vasilache committed Jan 14, 2020
1 parent 8d07f8d commit f52d717
Show file tree
Hide file tree
Showing 12 changed files with 458 additions and 312 deletions.
178 changes: 106 additions & 72 deletions mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
View file
Expand Up @@ -35,38 +35,9 @@ def StructuredOpTraits : NativeOpTrait<"linalg::StructuredOpTraits">;
// interface.
def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
let methods = [
InterfaceMethod<
"Query the number of inputs from the current operation.",
"unsigned", "getNumInputs"
>,
InterfaceMethod<
"Query the number of outputs from the current operation.",
"unsigned", "getNumOutputs"
>,
InterfaceMethod<
"Query the number of inputs and outputs from the current operation.",
"unsigned", "getNumInputsAndOutputs"
>,
InterfaceMethod<
"Query the input operands from the current operation.",
"Operation::operand_range", "getInputs"
>,
InterfaceMethod<
"Query the output operands from the current operation.",
"Operation::operand_range", "getOutputs"
>,
InterfaceMethod<
"Query the input and output operands from the current operation.",
"Operation::operand_range", "getInputsAndOutputs"
>,
InterfaceMethod<
"Query the iterator types attribute within the current operation.",
"ArrayAttr", "iterator_types"
>,
InterfaceMethod<
"Query the indexing maps attribute within the current operation.",
"ArrayAttr", "indexing_maps"
>,
//========================================================================//
// Loop types handling.
//========================================================================//
InterfaceMethod<
"Query the number of parallel loops within the current operation.",
"unsigned", "getNumParallelLoops"
Expand All @@ -82,40 +53,98 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
InterfaceMethod<
"Query the number of loops within the current operation.",
"unsigned", "getNumLoops">,

//========================================================================//
// Input arguments handling.
//========================================================================//
InterfaceMethod<
"Query the number of inputs from the current operation.",
"unsigned", "getNumInputs"
>,
InterfaceMethod<"Query the input view at the given index.",
"Value ", "getInput", (ins "unsigned":$i)
>,
InterfaceMethod<"Query the output view at the given index.",
"Value ", "getOutput", (ins "unsigned":$i)
>,
InterfaceMethod<[{
Return the index of the given input value `v`, or `None` if the value is
not an input.
}],
"llvm::Optional<unsigned>", "getIndexOfInput", (ins "Value ":$v)
>,
InterfaceMethod<[{
Query the index of the given view value, or `None` if the value is not
a view.
}],
"llvm::Optional<unsigned>", "getIndexOfOutput", (ins "Value ":$view)
InterfaceMethod<
"Query the input operands from the current operation.",
"Operation::operand_range", "getInputs"
>,
InterfaceMethod<[{
Query the type of the input shape at the given index.
}], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>,
InterfaceMethod<[{
Query the type of the output view at the given index.
}], "ShapedType", "getOutputShapedType", (ins "unsigned":$i)>,
InterfaceMethod<[{
Query whether the op has only MemRef input and outputs.
}], "bool", "hasBufferSemantics">,
InterfaceMethod<[{
Query the subset of input operands that are of ranked tensor type.
}], "SmallVector<RankedTensorType, 4>", "getInputTensorTypes">,


//========================================================================//
// Output arguments handling.
//========================================================================//
InterfaceMethod<
"Query the number of outputs from the current operation.",
"unsigned", "getNumOutputs"
>,
InterfaceMethod<"Query the output buffer at the given index.",
"Value ", "getOutputBuffer", (ins "unsigned":$i)
>,
InterfaceMethod<[{
Query the subset of output operands that are of ranked tensor type.
Query the index of the given buffer value, or `None` if the value is not
part of the output buffers.
}],
"llvm::Optional<unsigned>", "getIndexOfOutputBuffer", (ins "Value ":$view)
>,
InterfaceMethod<[{
Query the type of the output buffer at the given index.
}], "MemRefType", "getOutputBufferType", (ins "unsigned":$i)>,
InterfaceMethod<[{
Query the results that are of ranked tensor type.
}], "SmallVector<RankedTensorType, 4>", "getOutputTensorTypes">,
InterfaceMethod<
"Query the output buffers (operands) from the current operation.",
"Operation::operand_range", "getOutputBuffers"
>,

//========================================================================//
// Input and Output arguments handling.
//========================================================================//
InterfaceMethod<
"Return the number of inputs and outputs, irrespective of their buffer "
"or tensor type.",
"unsigned", "getNumInputsAndOutputs"
>,
InterfaceMethod<
"Return the number of inputs, irrespective of their buffer or tensor "
"type, and output buffers",
"unsigned", "getNumInputsAndOutputBuffers"
>,
InterfaceMethod<
"Return the range over inputs (irrespective of type) and output buffers.",
"Operation::operand_range", "getInputsAndOutputBuffers"
>,

//========================================================================//
// Other interface methods.
//========================================================================//
InterfaceMethod<
"Query the iterator types attribute within the current operation.",
"ArrayAttr", "iterator_types"
>,
InterfaceMethod<
"Query the indexing maps attribute within the current operation.",
"ArrayAttr", "indexing_maps"
>,
InterfaceMethod<[{
Query whether the op has only MemRef input and outputs.
}], "bool", "hasBufferSemantics">,

//========================================================================//
// Other static interface methods.
//========================================================================//
StaticInterfaceMethod<[{
Create an operation of the current type with the given location,
operands, and attributes.
Expand All @@ -128,9 +157,6 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
attributes);
}]
>,

/// Clone an operation with the given location and operands. This is used to
/// abstract away the optional underlying region creation.
InterfaceMethod<[{
Clone the current operation with the given location and operands. This
is used to abstract away the optional underlying region creation.
Expand Down Expand Up @@ -536,22 +562,26 @@ def GenericOp : GenericOpBase<"generic"> {
mixing input and output ranked tensor values with input and output memrefs.

```mlir
%1 = linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
%C = linalg.generic #trait_attribute %A, %B {other-attributes} :
tensor<?x?xf32>,
memref<?x?xf32, stride_specification>,
tensor<?x?xf32>
memref<?x?xf32, stride_specification>
-> (tensor<?x?xf32>)
```

In this case, the number of return values must match the number of output
tensor arguments. The semantics is that the `linalg.generic` op
produces (i.e. allocates and fills) its return values.
In this case, the number of outputs (args_out) must match the sum of (1) the
number of output buffer operands and (2) the number of tensor return values.
The semantics is that the `linalg.indexed_generic` op produces (i.e.
allocates and fills) its tensor return values.

Tensor values must be legalized by a buffer allocation pass before most
transformations can be applied. In particular, transformations that create
control flow around linalg.generic operations are not expected to mix with
tensors because SSA values do not escape naturally. Still, transformations
and rewrites that take advantage of tensor SSA values are expected to be
useful and will be added in the near future.
transformations can be applied. Such legalization moves tensor return values
into output buffer operands and updates the region arguments accordingly.

Transformations that create control-flow around linalg.indexed_generic
operations are not expected to work with tensors because SSA values do not
escape naturally. Still, transformations and rewrites that take advantage of
tensor SSA values are expected to be useful and will be added in the near
future.
}];
let verifier = [{ return ::verify(*this); }];
}
Expand Down Expand Up @@ -659,22 +689,26 @@ def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
memrefs.

```mlir
%1 = linalg.indexed_generic #trait_attribute %A, %B, %C {other-attributes}
%C = linalg.indexed_generic #trait_attribute %A, %B {other-attributes}
: tensor<?x?xf32>,
memref<?x?xf32, stride_specification>,
tensor<?x?xf32>
memref<?x?xf32, stride_specification>
-> (tensor<?x?xf32>)
```

In this case, the number of return values must match the number of output
tensor arguments. The semantics is that the `linalg.indexed_generic` op
produces (i.e. allocates and fills) its return values.
In this case, the number of outputs (args_out) must match the sum of (1) the
number of output buffer operands and (2) the number of tensor return values.
The semantics is that the `linalg.indexed_generic` op produces (i.e.
allocates and fills) its return values.

Tensor values must be legalized by a buffer allocation pass before most
transformations can be applied. In particular, transformations that create
control flow around linalg.generic operations are not expected to mix with
tensors because SSA values do not escape naturally. Still, transformations
and rewrites that take advantage of tensor SSA values are expected to be
useful and will be added in the near future.
transformations can be applied. Such legalization moves tensor return values
into output buffer operands and updates the region argument accordingly.

Transformations that create control-flow around linalg.indexed_generic
operations are not expected to work with tensors because SSA values do not
escape naturally. Still, transformations and rewrites that take advantage of
tensor SSA values are expected to be useful and will be added in the near
future.
}];
let verifier = [{ return ::verify(*this); }];
}
Expand Down

0 comments on commit f52d717

Please sign in to comment.