[Mlir-commits] [mlir] db9713c - [mlir] Add Tosa dialect const folder for tosa.const.

[Stella Laurenzo]

Author: Stella Laurenzo
Date: 2020-11-24T17:33:00Z
New Revision: db9713cd776ac5963efc502a5db6b315335aad9a

URL: https://github.com/llvm/llvm-project/commit/db9713cd776ac5963efc502a5db6b315335aad9a
DIFF: https://github.com/llvm/llvm-project/commit/db9713cd776ac5963efc502a5db6b315335aad9a.diff

LOG: [mlir] Add Tosa dialect const folder for tosa.const.

* Was missed in the initial submission and is required for a ConstantLike op.
* Also adds a materializeConstant hook to preserve it.
* Tightens up the argument constraint on tosa.const to match what is actually legal.

Differential Revision: https://reviews.llvm.org/D92040

Added: 
    mlir/test/Dialect/Tosa/constant_folding.mlir

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
    mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
    mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
index 5701e6ec97b3..2d977d76df9f 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
@@ -38,6 +38,7 @@ def Tosa_Dialect : Dialect {
   }];
 
   let cppNamespace = "mlir::tosa";
+  let hasConstantMaterializer = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 43e8bfacee27..c9790596ed88 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1512,12 +1512,13 @@ def Tosa_ConstOp : Tosa_Op<"const", [ConstantLike, NoSideEffect,
   }];
 
   let arguments = (ins
-    AnyAttr:$value
+    ElementsAttr:$value
   );
 
   let results = (outs
     Tosa_TensorUpto4D:$output
   );
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index daf10b2013bf..a609e64cb7d7 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -89,6 +89,24 @@ void TosaDialect::initialize() {
   addInterfaces<TosaInlinerInterface>();
 }
 
+Operation *TosaDialect::materializeConstant(OpBuilder &builder, Attribute value,
+                                            Type type, Location loc) {
+  // Tosa dialect constants only support ElementsAttr unlike standard dialect
+  // constant which supports all attributes.
+  if (value.isa<ElementsAttr>())
+    return builder.create<tosa::ConstOp>(loc, type, value.cast<ElementsAttr>());
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
+// Operator Folders.
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.empty() && "constant has no operands");
+  return valueAttr();
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator Verifiers.
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
new file mode 100644
index 000000000000..988de5c7a9ac
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -0,0 +1,8 @@
+// RUN: mlir-opt --test-constant-fold %s | FileCheck %s
+
+// CHECK-LABEL: func @test_const
+func @test_const(%arg0 : index) -> tensor<4xi32> {
+  // CHECK: "tosa.const"
+  %0 = "tosa.const"() {value = dense<[3, 0, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
+  return %0 : tensor<4xi32>
+}