Page MenuHomePhabricator

[mlir][tosa] Support dynamic batch dimension for ops where the batch dim is explicit
AcceptedPublic

Authored by NatashaKnk on Dec 6 2021, 3:46 PM.

Details

Reviewers
hanchung
sjarus
Summary

Conv2D, DepthwiseConv2D, MaxPool2D, AvgPool2D, Rescale, Resize ang Gather lowering batch support for tosa-to-linalg lowering

Diff Detail

Event Timeline

NatashaKnk created this revision.Dec 6 2021, 3:46 PM
NatashaKnk requested review of this revision.Dec 6 2021, 3:46 PM
hanchung requested changes to this revision.Dec 7 2021, 10:49 AM
hanchung added inline comments.
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
1056

There are assumptions in the method. It is not a simple method, so we should have function documentation.

E.g., All the dims other than batch dim is expected to be static, the batch dim is dim #0, etc.

1056–1058

We prefer returning LogicalResult in LLVM. In this case, you can return success() or failure or rewriter.noitfyMatchFailure(...).

After reading a bit more, I think it's better to return Optional<SmallVector<Value>>. You can return llvm::None for failure cases.

1058

s/SmallVector<Value>/SmallVectorImpl<Value>

1061

Use auto because cast already spells the type.

https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

Don’t “almost always” use auto, but do use auto with initializers like cast<Foo>(...) or other places where the type is already obvious from the context.

1062–1064
1072–1073

I think you can write it simpler and more LLVM style like

if (llvm::any(dynTy.getShape().drop_front(), ShapedType::isDynamic)) {
  ...
}
1081–1082

Should we iterate through all params? Do all the params share the same dynamic size?

1128

Doesn't op.output() work? `It's more readable and people would know that there is only single result for the op.

1309

ditto

1989–1992

do not use braces for simple single if-check.

This revision now requires changes to proceed.Dec 7 2021, 10:49 AM

Add more llvm-style formatting, changed return type of checkForDynamicBatchOnly function

hanchung accepted this revision.Dec 16 2021, 5:44 PM
hanchung added inline comments.
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
1993

Delete empty statement?

1999–2000

dynamicDims is only used once, we can write it this way.

auto dynamicDims =
        checkForDynamicBatchOnly(rewriter, op, {input, op.output()});
if (!dynamicDims.hasValue())
      return failure();
...
create<linalg::InitTensorOp>(loc, dynamicDims.getValue(), resultTy.getShape(),
...

Same for other places

This revision is now accepted and ready to land.Dec 16 2021, 5:44 PM