From df6003ff321a419a5cbd9bcd5b8641936d22f6a2 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell@arm.com>
Date: Thu, 29 Feb 2024 15:20:34 +0000
Subject: [PATCH] [mlir][VectorOps][nfc] Add result pretty printing to
 `vector.vscale`

This will now print the value of `vector.vscale` as `%vscale` in IR
dumps which makes it easier to spot where things are scalable.

One test that depended on the value names has also been fixed.
---
 mlir/include/mlir/Dialect/Vector/IR/VectorOps.td       | 10 +++++++++-
 .../Linalg/transform-op-peel-and-vectorize.mlir        |  2 +-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 6d50b0654bc57..06360bd10e525 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2809,7 +2809,8 @@ def Vector_SplatOp : Vector_Op<"splat", [
 //       call to the function. For that, it might be useful to have a
 //       'vector.scale.global' and a 'vector.scale.local' operation.
 def VectorScaleOp : Vector_Op<"vscale",
-                 [Pure]> {
+  [Pure, DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]
+> {
   let summary = "Load vector scale size";
   let description = [{
     The `vscale` op returns the scale of the scalable vectors, a positive
@@ -2825,6 +2826,13 @@ def VectorScaleOp : Vector_Op<"vscale",
   }];
   let results = (outs Index:$res);
   let assemblyFormat = "attr-dict";
+
+  let extraClassDefinition = [{
+    void $cppClass::getAsmResultNames(
+        ::llvm::function_ref<void(mlir::Value, mlir::StringRef)> setNameFn) {
+      setNameFn(getResult(), "vscale");
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
index d54cace31efb9..b7e316f8925d3 100644
--- a/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-peel-and-vectorize.mlir
@@ -19,7 +19,7 @@ func.func @matmul(%A: tensor<1024x512xf32>,
 // CHECK-DAG:  %[[C16:.*]] = arith.constant 16 : index
 // CHECK:      %[[VSCALE:.*]] = vector.vscale
 // CHECK:      %[[STEP:.*]] = arith.muli %[[VSCALE]], %[[C16]] : index
-// CHECK:      %2 = scf.for {{.*}} %[[C0]] to %[[C1024]] step %[[C8]] iter_args(%arg4 = %arg2) -> (tensor<1024x2000xf32>) {
+// CHECK:      scf.for {{.*}} %[[C0]] to %[[C1024]] step %[[C8]] iter_args(%arg4 = %arg2) -> (tensor<1024x2000xf32>) {
 
 // Main loop after vectorisation (without masking)