[mlir][tosa] Enable `apply_scale` unrolling

[kuhar]

Summary

Make tosa.apply_scale implement VectorUnrollOpInterface so that we
can unroll it.

[kuhar]

I didn't find an easy way to test this. The closest existing thing I could find is -test-vector-unrolling-patterns but it seems to only work on a number of hardcoded ops.

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

 #ifndef MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 #define MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/VectorInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // TOSA dialect and structs includes.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.h.inc"
 
 namespace mlir {

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

 
 #ifndef TOSA_UTIL_OPS
 #define TOSA_UTIL_OPS
 
 include "mlir/IR/OpBase.td"
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
+include "mlir/Interfaces/VectorInterfaces.td"
 include "mlir/Dialect/Tosa/IR/TosaInterfaces.td"
 
 include "mlir/Dialect/Tosa/IR/TosaTypesBase.td"
 include "mlir/Dialect/Tosa/IR/TosaOpBase.td"
 
-def Tosa_ApplyScaleOp: Tosa_Op<"apply_scale", [Pure] # ElementwiseMappable.traits> {
+def Tosa_ApplyScaleOp :
+  Tosa_Op<"apply_scale",
+          [Pure, DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+          ElementwiseMappable.traits> {
   let summary = "Rescale scalar operator for Tosa tensor operators";
 
   let description = [{
     Applies rescaling for fixed point values. This behavior is replicated in
     multiple quantized operations (mul, convolution, rescale, matmul, pooling).
 
     The commonplace implementation is to use i64 operations to avoid integer
     overflow with target specific implementations can use native operations to
     avoid wider than necessary types.
   }];
 
   let arguments = (ins
     Tosa_IntLike:$value,
     Tosa_IntLike:$multiplier,
     Tosa_Int8Like:$shift,
     BoolAttr:$double_round
   );
 
   let results = (outs
     Tosa_IntLike:$output
   );
+
+  let extraClassDeclaration = [{
+    std::optional<SmallVector<int64_t, 4>> getShapeForUnroll();
+  }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: yield
 //===----------------------------------------------------------------------===//
 def Tosa_YieldOp : Tosa_Op<"yield", [
        Terminator,
        Pure]> {

mlir/lib/Dialect/Tosa/CMakeLists.txt

add_mlir_dialect_library(MLIRTosaDialect
   MLIRIR
   MLIRDialect
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRQuantDialect
   MLIRQuantUtils
   MLIRSideEffectInterfaces
   MLIRTensorDialect
+  MLIRVectorInterfaces
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Transforms)

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

LogicalResult WhileOp::inferReturnTypeComponents(
 
   for (const ValueKnowledge &result : resultKnowledge) {
     inferredReturnShapes.push_back(result.getShapedTypeComponents());
   }
 
   return success();
 }
 
+std::optional<SmallVector<int64_t, 4>> ApplyScaleOp::getShapeForUnroll() {
+  if (auto vt = getType().dyn_cast<VectorType>())
+    return llvm::to_vector<4>(vt.getShape());
+  return std::nullopt;
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Attribute Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 td_library(
     name = "TosaDialectTdFiles",
     srcs = glob(["include/mlir/Dialect/Tosa/IR/*.td"]),
     deps = [
         ":InferTypeOpInterfaceTdFiles",
         ":LoopLikeInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
+        ":VectorInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "TosaDialectIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (