This is an archive of the discontinued LLVM Phabricator instance.

Differential D79764

[mlir][Linalg] Allow reshapes to collapse to a zero-rank tensor.
ClosedPublic

Authored by mravishankar on May 11 2020, 10:51 PM.

Details

Reviewers
nicolasvasilache
Commits
rGd2a956985016: [mlir][Linalg] Allow reshapes to collapse to a zero-rank tensor.
Summary

This is only valid if the source tensors (result tensor) is static
shaped with all unit-extents when the reshape is collapsing
(expanding) dimensions.

Depends On D79763

Diff Detail

Event Timeline

mravishankar created this revision.May 11 2020, 10:51 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptMay 11 2020, 10:51 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: llvm-commits, Kayjukh, frgossen and 14 others. · View Herald Transcript
mravishankar added a child revision: D79765: [mlir][Linalg] Add folders and canonicalizers for linalg.reshape/linalg.tensor_reshape operations..May 11 2020, 10:52 PM
Harbormaster completed remote builds in B56399: Diff 263349.May 11 2020, 11:57 PM
nicolasvasilache accepted this revision.May 12 2020, 5:45 PM

Looks good but please update the doc as appropriate.

This revision is now accepted and ready to land.May 12 2020, 5:45 PM
mravishankar updated this revision to Diff 263608.May 12 2020, 10:24 PM

Addressing comments

Harbormaster completed remote builds in B56542: Diff 263608.May 12 2020, 11:09 PM
Closed by commit rGd2a956985016: [mlir][Linalg] Allow reshapes to collapse to a zero-rank tensor. (authored by mravishankar). · Explain WhyMay 12 2020, 11:09 PM
This revision was automatically updated to reflect the committed changes.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
IR/
12 lines
test/
Dialect/
Linalg/
29 lines
23 lines

Diff 263349

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

Show First 20 Lines • Show All 547 Lines • ▼ Show 20 Linesstatic LogicalResult verifyReshapeLikeTypes(Op op, T &expandedType,
collapsedType = op.getResultType(); collapsedType = op.getResultType();
unsigned expandedRank = expandedType.getRank(); unsigned expandedRank = expandedType.getRank();
unsigned collapsedRank = collapsedType.getRank(); unsigned collapsedRank = collapsedType.getRank();
bool isCollapse = expandedRank > collapsedRank; bool isCollapse = expandedRank > collapsedRank;
if (!isCollapse) { if (!isCollapse) {
std::swap(expandedRank, collapsedRank); std::swap(expandedRank, collapsedRank);
std::swap(expandedType, collapsedType); std::swap(expandedType, collapsedType);
} }
if (expandedRank == 0 || collapsedRank == 0) if (expandedRank == 0)
return op.emitOpError("expected non-zero memref ranks"); return op.emitOpError("expected non-zero memref ranks");
if (expandedRank == collapsedRank) if (expandedRank == collapsedRank)
return op.emitOpError("expected to collapse or expand dims"); return op.emitOpError("expected to collapse or expand dims");
if (collapsedRank == 0) {
// If collapsed rank is 0, then expanded type must be static shaped and of
// sizes 1.
if (llvm::any_of(expandedType.getShape(),
[](int64_t dim) -> bool { return dim != 1; }))
return op.emitOpError(
"invalid to reshape tensor/memref with non-unit extent dimensions to "
"zero-rank tensor/memref");
return success();
}
if (collapsedRank != op.reassociation().size()) if (collapsedRank != op.reassociation().size())
return op.emitOpError("expected rank of the collapsed type(") return op.emitOpError("expected rank of the collapsed type(")
<< collapsedRank << ") to be the number of reassociation maps(" << collapsedRank << ") to be the number of reassociation maps("
<< op.reassociation().size() << ")"; << op.reassociation().size() << ")";
auto maps = getAffineMaps(op.reassociation()); auto maps = getAffineMaps(op.reassociation());
for (auto it : llvm::enumerate(maps)) for (auto it : llvm::enumerate(maps))
if (it.value().getNumDims() != expandedRank) if (it.value().getNumDims() != expandedRank)
return op.emitOpError("expected reassociation map #") return op.emitOpError("expected reassociation map #")
▲ Show 20 Lines • Show All 653 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/llvm.mlir

Show First 20 Lines • Show All 267 Lines • ▼ Show 20 Lines
// CHECK: llvm.mlir.constant(5 : index) : !llvm.i64 // CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
// CHECK: llvm.insertvalue {{.*}}[3, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> // CHECK: llvm.insertvalue {{.*}}[3, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
// CHECK: llvm.mlir.constant(20 : index) : !llvm.i64 // CHECK: llvm.mlir.constant(20 : index) : !llvm.i64
// CHECK: llvm.insertvalue {{.*}}[4, 0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> // CHECK: llvm.insertvalue {{.*}}[4, 0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
// CHECK: llvm.mlir.constant(5 : index) : !llvm.i64 // CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
// CHECK: llvm.insertvalue {{.*}}[4, 1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> // CHECK: llvm.insertvalue {{.*}}[4, 1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
// CHECK: llvm.mlir.constant(1 : index) : !llvm.i64 // CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
// CHECK: llvm.insertvalue {{.*}}[4, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> // CHECK: llvm.insertvalue {{.*}}[4, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
func @reshape_zero_dim(%arg0 : memref<1x1xf32>) {
%0 = linalg.reshape %arg0 [] : memref<1x1xf32> into memref<f32>
%1 = linalg.reshape %0 [] : memref<f32> into memref<1x1xf32>
return
}
// CHECK-LABEL: func @reshape_zero_dim
// CHECK: llvm.mlir.undef : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.extractvalue %{{.*}}[0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.extractvalue %{{.*}}[1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.extractvalue %{{.*}}[2] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.mlir.undef : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.extractvalue %{{.*}}[0] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.extractvalue %{{.*}}[1] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.extractvalue %{{.*}}[2] : !llvm<"{ float*, float*, i64 }">
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
// CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
// CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[4, 1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">

mlir/test/Dialect/Linalg/roundtrip.mlir

Show First 20 Lines • Show All 627 Lines • ▼ Show 20 Linesfunc @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,
linalg.batch_matmul %a3, %b3, %c3 : (memref<?x?x?xf32>, memref<?x?x?xf32>, memref<?x?x?xf32>) -> () linalg.batch_matmul %a3, %b3, %c3 : (memref<?x?x?xf32>, memref<?x?x?xf32>, memref<?x?x?xf32>) -> ()
linalg.batch_matmul %ta3, %tb3, %c3 : (tensor<?x?x?xf32>, tensor<?x?x?xf32>, memref<?x?x?xf32>) -> () linalg.batch_matmul %ta3, %tb3, %c3 : (tensor<?x?x?xf32>, tensor<?x?x?xf32>, memref<?x?x?xf32>) -> ()
return return
} }
// CHECK-LABEL: func @named_ops // CHECK-LABEL: func @named_ops
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// -----
func @tensor_reshape_zero_dim(%arg0 : tensor<1x1xf32>, %arg1 : tensor<f32>) -> (tensor<f32>, tensor<1x1xf32>)
{
%0 = linalg.tensor_reshape %arg0 [] : tensor<1x1xf32> into tensor<f32>
%1 = linalg.tensor_reshape %0 [] : tensor<f32> into tensor<1x1xf32>
return %0, %1 : tensor<f32>, tensor<1x1xf32>
}
// CHECK-LABEL: func @tensor_reshape_zero_dim
// CHECK: linalg.tensor_reshape %{{.*}} [] : tensor<1x1xf32> into tensor<f32>
// CHECK: linalg.tensor_reshape %{{.*}} [] : tensor<f32> into tensor<1x1xf32>
// -----
func @memref_reshape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>) -> (memref<f32>, memref<1x1xf32>)
{
%0 = linalg.reshape %arg0 [] : memref<1x1xf32> into memref<f32>
%1 = linalg.reshape %0 [] : memref<f32> into memref<1x1xf32>
return %0, %1 : memref<f32>, memref<1x1xf32>
}
// CHECK-LABEL: func @memref_reshape_zero_dim
// CHECK: linalg.reshape %{{.*}} [] : memref<1x1xf32> into memref<f32>
// CHECK: linalg.reshape %{{.*}} [] : memref<f32> into memref<1x1xf32>