[MLIR][TOSA] Added Tosa to Standard/SCF Lowerings (const, if, while)

[rsuderman]

Summary

Includes a lowering for tosa.const, tosa.if, and tosa.while to Standard/SCF dialects. TosaToStandard is
used for constant lowerings and TosaToSCF handles the if/while ops.

[rsuderman]

Fixed linter errors.

[silvas]

test cases?

[rsuderman]

Fixed upload

[silvas]

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

48 ↗

(On Diff #326164)

it feels like these could be in TosaToSCF and convert to scf.if/scf.while. Would be easier to write them as patterns too.

[rriddle]

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

77

Remove all of the mlir:: in this commit.

[rriddle]

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

48 ↗

(On Diff #326164)

+1

[rsuderman]

Moved control flow lowerings to a TosaToSCF set of passes.

[rsuderman]

Fixed arc lint error.

[rsuderman]

clang-tidy error

[silvas]

mlir/include/mlir/Conversion/TosaToStandard/TosaToStandard.h

23 ↗

(On Diff #326459)

decl not needed?

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

2

nit: formatting

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

31 ↗

(On Diff #326459)

specify in td file.

38 ↗

(On Diff #326459)

you should only need ConstantOp legal.

45 ↗

(On Diff #326459)

use partial conversion so you don't need markUnknownOpDynamicallyLegal

mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir

39 ↗

(On Diff #326459)

nit: indent

[rriddle]

mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp

26 ↗

(On Diff #326459)

Functions should be marked static and at the top-level, only classes should be in anonymous namespaces.

30 ↗

(On Diff #326459)

nit: Spell out auto here.

31 ↗

(On Diff #326459)

Remove trivial braces.

42 ↗

(On Diff #326459)

This seems sketchy, how does this work with rollbacks? Why not rewriter.erase?

66 ↗

(On Diff #326459)

Same comment here.

mlir/lib/Conversion/TosaToSCF/TosaToSCFPass.cpp

29 ↗

(On Diff #326459)

Why does this need to be a FuncOp pass?

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

35 ↗

(On Diff #326459)

Same comment, why is this a FuncOp pass?

[rsuderman]

silvas@ and rriddle@ comments

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

 //===- TosaToStandard.cpp - Lowering Tosa to Standard Dialect
 //-----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // These rewriters lower from the Tosa to the Standard dialect.
LogicalResult matchAndRewrite(mlir::tosa::ConstOp op,
                                 PatternRewriter &rewriter) const final {
     rewriter.replaceOpWithNewOp<mlir::ConstantOp>(op, op.value());
     return success();
   }
 };
 
 // Replaces all tosa.yield ops with a branch to the correct target. This is
 // used for processing each branch of tosa.if
-LogicalResult ReplaceYieldOps(Region *region, Block *targetBlock, Location loc,
+LogicalResult replaceYieldOps(Region *region, Block *targetBlock, Location loc,
                               const BlockAndValueMapping &mapper,
                               OpBuilder *builder) {
   for (auto &oldBlock : region->getBlocks()) {
     Block *block = mapper.lookup(&oldBlock);
     auto returnOp = dyn_cast<tosa::YieldOp>(block->getTerminator());
     if (!returnOp)
       continue;
     builder->setInsertionPointToEnd(block);
     builder->create<mlir::BranchOp>(loc, targetBlock, returnOp.getOperands());
     returnOp.erase();
   }
 
   return success();
 }
 
 LogicalResult legalizeTosaIf(IfOp ifOp) {
   Operation *operation = ifOp.getOperation();
   mlir::OpBuilder builder(ifOp);
-  auto origBlock = operation->getBlock();
+  auto *origBlock = operation->getBlock();
   auto *tailBlock = origBlock->splitBlock(operation);
   auto loc = ifOp.getLoc();
 
   // Initially we clone the then/else blocks at the current point in the
   // region so we can branch to them in the conditional check.
   BlockAndValueMapping mapper;
   ifOp.then_branch().cloneInto(origBlock->getParent(),
                                Region::iterator(tailBlock), mapper);
   ifOp.else_branch().cloneInto(origBlock->getParent(),
                                Region::iterator(tailBlock), mapper);
 
   Block *thenBranch = mapper.lookup(&ifOp.then_branch().front());
   Block *elseBranch = mapper.lookup(&ifOp.else_branch().front());
 
   // Move to the conditional section, extract the value from the condition
   // tensor, and execute the conditional branch to the then/else cases.
   builder.setInsertionPointToEnd(origBlock);
   auto condValue = builder.create<mlir::tensor::ExtractOp>(loc, ifOp.cond());
   builder.create<mlir::CondBranchOp>(loc, condValue, thenBranch, ifOp.inputs(),
                                      elseBranch, ifOp.inputs());
 
   // Yields need to be replaced with branches to the tile block.
-  if (failed(ReplaceYieldOps(&ifOp.then_branch(), tailBlock, loc, mapper,
+  if (failed(replaceYieldOps(&ifOp.then_branch(), tailBlock, loc, mapper,
                              &builder)))
     return failure();
-  if (failed(ReplaceYieldOps(&ifOp.else_branch(), tailBlock, loc, mapper,
+  if (failed(replaceYieldOps(&ifOp.else_branch(), tailBlock, loc, mapper,
                              &builder)))
     return failure();
 
   tailBlock->addArguments(ifOp.getResultTypes());
 
   for (auto it : llvm::enumerate(ifOp.getResults())) {
     it.value().replaceAllUsesWith(tailBlock->getArgument(it.index()));
   }
LogicalResult legalizeTosaWhile(WhileOp whileOp) {
   // Move back to before the cond/body blocks and branch to the conditiono op.
   builder.setInsertionPointToEnd(origBlock);
   builder.create<mlir::BranchOp>(loc, condBlock, whileOp.getOperands());
 
   // Update the inlined condition blocks by replacing any yield ops with an
   // tensor.extract and conditional branch to the body block or tail.
   builder.setInsertionPointToStart(condBlock);
   for (auto &block : whileOp.cond()) {
-    auto newBlock = mapper.lookup(&block);
+    auto *newBlock = mapper.lookup(&block);
 
-    auto returnOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
-    if (!returnOp)
+    auto yieldOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
+    if (!yieldOp)
       continue;
     builder.setInsertionPointToEnd(newBlock);
 
-    auto return_value = returnOp.getOperand(0);
-    auto condValue = builder.create<mlir::tensor::ExtractOp>(loc, return_value);
+    auto yieldValue = yieldOp.getOperand(0);
+    auto condValue = builder.create<mlir::tensor::ExtractOp>(loc, yieldValue);
 
     // Pass the args passed to the cond block as the arguments for each branch
     // condition.
     llvm::SmallVector<Value, 4> successorArgs(condBlock->args_begin(),
                                               condBlock->args_end());
     builder.create<mlir::CondBranchOp>(loc, condValue, bodyBlock, successorArgs,
                                        tailBlock, successorArgs);
-    returnOp.erase();
+    yieldOp.erase();
   }
 
   // Replaces all yield s in the while op body with branches to the condition.
   for (auto &block : whileOp.body()) {
-    auto newBlock = mapper.lookup(&block);
-    auto returnOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
-    if (!returnOp)
+    auto *newBlock = mapper.lookup(&block);
+    auto yieldOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
+    if (!yieldOp)
       continue;
     builder.setInsertionPointToEnd(newBlock);
-    builder.create<mlir::BranchOp>(loc, condBlock, returnOp.getOperands());
-    returnOp.erase();
+    builder.create<mlir::BranchOp>(loc, condBlock, yieldOp.getOperands());
+    yieldOp.erase();
   }
 
   // Erase the original while loop.
   tailBlock->addArguments(whileOp.getResultTypes());
   for (auto it : llvm::enumerate(whileOp.getResults())) {
     it.value().replaceAllUsesWith(tailBlock->getArgument(it.index()));
   }
   operation->erase();