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

[stellaraccident]

Summary

• 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.

mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td

let description = [{
     implementations. Applications or frameworks which target TOSA can therefore
     be deployed on a wide range of different processors, such as CPUs or GPUs,
     with defined accuracy and compatibility constraints. Most operators from the
     common ML frameworks should be expressible in TOSA. It is expected that
     there will be tools to lower from the ML frameworks into TOSA.
   }];
 
   let cppNamespace = "mlir::tosa";
+  let hasConstantMaterializer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Quantization Attributes.
 //===----------------------------------------------------------------------===//
 
 // Quantization attributes used across TOSA operators. Quantization attributes
 // feed numerical precision parameters to the functional implementation of TOSA

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_ConstOp : Tosa_Op<"const", [ConstantLike, NoSideEffect,
   let summary = "Constant op.";
 
   let description = [{
     A node containing constant data for use as the input to an operation. May
     hold data in any of the supported data formats.
   }];
 
   let arguments = (ins
-    AnyAttr:$value
+    ElementsAttr:$value
   );
 
   let results = (outs
     Tosa_TensorUpto4D:$output
   );
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: identity
 //===----------------------------------------------------------------------===//
 def Tosa_IdentityOp: Tosa_Op<"identity", [NoSideEffect]> {
   let summary = "Identity operator";
   let description = [{

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 void TosaDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
       >();
   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.
 //===----------------------------------------------------------------------===//
 
 template <typename T>
 static LogicalResult verifyConvOp(T op) {
   // All TOSA conv ops have an input() and weight().
   auto inputType = op.input().getType().template dyn_cast<RankedTensorType>();

mlir/test/Dialect/Tosa/constant_folding.mlir

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