[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

@@ -38,6 +38,7 @@ def Tosa_Dialect : Dialect {
   }];
 
   let cppNamespace = "mlir::tosa";
+  let hasConstantMaterializer = 1;
 }
 
 //===----------------------------------------------------------------------===//

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;
 }
 
 //===----------------------------------------------------------------------===//

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.
 //===----------------------------------------------------------------------===//

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>
+}