maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[MLIR] Add missing traits and assembly format to `shape.from/to_extent_tensor` 

Add `NoSideEffect` trait to `shape.to_extent_tensor` and
`shape.from_extent_tensor` and defined custom assembly format for the
operations.

Differential Revision: https://reviews.llvm.org/D81158
This commit is contained in:
Frederik Gossen 2020-06-12 10:53:15 +00:00
parent 4495a6b141
commit 6196c37969
3 changed files with 18 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
View File

@ -155,7 +155,7 @@ def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
let hasFolder = 1; let hasFolder = 1;
} }
def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> { def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
let summary = "Creates a shape from a tensor of extents"; let summary = "Creates a shape from a tensor of extents";
let description = [{ let description = [{
Creates a shape from a 1D integral tensor of extents. The rank of the Creates a shape from a 1D integral tensor of extents. The rank of the
@ -165,26 +165,25 @@ def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
let arguments = (ins IndexTensor:$input); let arguments = (ins IndexTensor:$input);
let results = (outs Shape_ShapeType:$result); let results = (outs Shape_ShapeType:$result);
let assemblyFormat = "attr-dict $input `:` type($input)";
} }
def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> { def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
let summary = "Creates a dimension tensor from a shape"; let summary = "Creates a dimension tensor from a shape";
// TODO: Think more about the error situation. Perhaps factor out the
// error detection into a separate op so downstream consumers can control
// their error behavior. Then this op would assume that the input has
// been properly checked to not be an error (and could thus be a
// NoSideEffect op).
let description = [{ let description = [{
Converts a shape to a 1D integral tensor of extents. The number of elements Converts a shape to a 1D integral tensor of extents. The number of elements
in the tensor equals the rank of the shape, and the elements equal the in the tensor equals the rank of the shape, and the elements equal the
extents of the shape. extents of the shape.
If the shape represents an error, then this op currently aborts the program. If the shape represents an error, this op's behavior is undefined.
}]; }];
let arguments = (ins Shape_ShapeType:$input); let arguments = (ins Shape_ShapeType:$input);
let results = (outs IndexTensor:$result); let results = (outs IndexTensor:$result);
let assemblyFormat = "attr-dict $input `:` type($result)";
let hasFolder = 1; let hasFolder = 1;
} }

2
mlir/test/Dialect/Shape/canonicalize.mlir
View File

@ -89,7 +89,7 @@ func @f() -> !shape.shape {
func @f() -> tensor<2xindex> { func @f() -> tensor<2xindex> {
// CHECK: constant dense<[0, 1]> : tensor<2xindex> // CHECK: constant dense<[0, 1]> : tensor<2xindex>
%cs = shape.const_shape [0, 1] %cs = shape.const_shape [0, 1]
%0 = "shape.to_extent_tensor"(%cs) : (!shape.shape) -> tensor<2xindex> %0 = shape.to_extent_tensor %cs : tensor<2xindex>
return %0 : tensor<2xindex> return %0 : tensor<2xindex>
} }

10
mlir/test/Dialect/Shape/ops.mlir
View File

@ -101,3 +101,13 @@ func @const_size() {
%2 = shape.const_size 2 %2 = shape.const_size 2
return return
} }
func @test_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
%0 = shape.to_extent_tensor %arg : tensor<3xindex>
return %0 : tensor<3xindex>
}
func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
%0 = shape.from_extent_tensor %arg : tensor<?xindex>
return %0 : !shape.shape
}