From db9713cd776ac5963efc502a5db6b315335aad9a Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Tue, 24 Nov 2020 16:58:26 +0000
Subject: [PATCH] [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
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td |  1 +
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td    |  3 ++-
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp            | 18 ++++++++++++++++++
 mlir/test/Dialect/Tosa/constant_folding.mlir    |  8 ++++++++
 4 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Tosa/constant_folding.mlir

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
index 5701e6e..2d977d7 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 43e8bfa..c979059 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 daf10b2..a609e64 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 0000000..988de5c
--- /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>
+}
-- 
2.7.4