This is an archive of the discontinued LLVM Phabricator instance.

Differential D82644

[MLIR][Shape] Lower `shape.get_extent` to `std.dim` when possible
ClosedPublic

Authored by frgossen on Jun 26 2020, 5:03 AM.

Details

Reviewers
jpienaar
herhut
silvas
Commits
rG76d72c941df1: [MLIR][Shape] Lower `shape.get_extent` to `std.dim` when possible
Summary

When the shape is derived from a tensor argument the shape extent can be derived
directly from that tensor with std.dim.
This lowering pattern circumvents the necessity to materialize the shape in
memory.

Diff Detail

Event Timeline

frgossen created this revision.Jun 26 2020, 5:03 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 26 2020, 5:03 AM
Herald added subscribers: msifontes, jurahul, Kayjukh and 14 others. · View Herald Transcript
frgossen added a child revision: D82645: [MLIR][Shape] Lower `shape.get_extent` to `extract_element` when possible.Jun 26 2020, 5:03 AM
herhut accepted this revision.Jun 26 2020, 5:13 AM

Neat!

This revision is now accepted and ready to land.Jun 26 2020, 5:13 AM
Harbormaster failed remote builds in B61910: Diff 273679!Jun 26 2020, 6:30 AM
Closed by commit rG76d72c941df1: [MLIR][Shape] Lower `shape.get_extent` to `std.dim` when possible (authored by frgossen). · Explain WhyJun 29 2020, 2:07 AM
This revision was automatically updated to reflect the committed changes.

Revision Contents

PathSize
mlir/
lib/
Conversion/
ShapeToStandard/
7 lines
test/
Conversion/
ShapeToStandard/
16 lines

Diff 273679

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

Show All 13 Lines
// represented as indices. // represented as indices.
def IndexToSizeOpConversion : Pat< def IndexToSizeOpConversion : Pat<
(Shape_IndexToSizeOp $arg), (Shape_IndexToSizeOp $arg),
(replaceWithValue $arg)>; (replaceWithValue $arg)>;
def SizeToIndexOpConversion : Pat< def SizeToIndexOpConversion : Pat<
(Shape_SizeToIndexOp $arg), (Shape_SizeToIndexOp $arg),
(replaceWithValue $arg)>; (replaceWithValue $arg)>;
// Derive shape extent directly from shape origin if possible.
// This circumvents the necessity to materialize the shape in memory.
def GetExtentShapeOfConversion : Pat<
(Shape_GetExtentOp (Shape_ShapeOfOp $arg), $idx),
(Shape_IndexToSizeOp (DimOp $arg, (Shape_SizeToIndexOp $idx))),
[],
(addBenefit 10)>;

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

Show First 20 Lines • Show All 121 Lines • ▼ Show 20 Lines
// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
func @rank(%shape : !shape.shape) -> !shape.size { func @rank(%shape : !shape.shape) -> !shape.size {
// CHECK-DAG: %[[C0:.*]] = constant 0 : index // CHECK-DAG: %[[C0:.*]] = constant 0 : index
// CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]] // CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
// CHECK-DAG: return %[[RESULT]] : index // CHECK-DAG: return %[[RESULT]] : index
%rank = shape.rank %shape %rank = shape.rank %shape
return %rank : !shape.size return %rank : !shape.size
} }
// -----
// Express `get_extent` as `std.dim` when it relies directly on the outcome of a
// `shape_of` operation.
// CHECK-LABEL: @get_extent_shape_of
// CHECK-SAME: (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : !shape.size)
-> !shape.size {
// CHECK: %[[RESULT:.*]] = dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
// CHECK: return %[[RESULT]] : index
%shape = shape.shape_of %arg : tensor<2x3xf32>
%result = shape.get_extent %shape, %idx
return %result : !shape.size
}