Remove redundant attributes from Rock ops (#1910)

* Remove arch and num_cu

* Remove comment

* Fallback to checking mod if arch not found on func

* Fall back to arch info for minNumCU

* Small fixes for LIT tests

* Update LIT tests for removal of arch and num_cu

* Remove additional whitespace

* Initial changes for making the features attribute optional

* Fix LIT errors with optional features attribute

* Update determination of types

* More fixes uncovered during LIT tests

* Updates to remove 'features' where possible from LIT tests

* Fix for review comments

* Add recursive type check to AmdArchDb.cpp

* More review changes

* Create GetRockInfo header and add 'features' attribute to funcs

* Fix compile issues

* More fixes for LIT tests

* Remove features from unnecessary ops

* Update additional LIT tests to remove 'features'

* Clang format

* Fix minor nitpicks

* clang-format

* Add 'getGemmFeaturesAttr()'

* Update assert to check for all features being the same

* clang-format AffixTuningParameters

* Move templated code to its own header

* Investigate tuna script issue

* Fixes for running out of heap space on Windows

* Revert "Investigate tuna script issue"

This reverts commit 9638f09.

* Fix small typos

Author: justinrosner

mlir/include/mlir/Dialect/Rock/utility/AmdArchDb.h renamed to mlir/include/mlir/Dialect/Rock/IR/AmdArchDb.h

@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_ROCK_UTILITY_AMDARCHDB_H
-#define MLIR_DIALECT_ROCK_UTILITY_AMDARCHDB_H
+#ifndef MLIR_DIALECT_ROCK_IR_AMDARCHDB_H
+#define MLIR_DIALECT_ROCK_IR_AMDARCHDB_H
 
 #include "mlir/Dialect/Rock/IR/Rock.h"
 #include "mlir/Support/LLVM.h"
@@ -51,4 +51,4 @@ AmdArchInfo lookupArchInfo(StringRef arch);
 } // namespace rock
 } // namespace mlir
 
-#endif // MLIR_DIALECT_ROCK_UTILITY_AMDARCHDB_H
+#endif // MLIR_DIALECT_ROCK_IR_AMDARCHDB_H

mlir/include/mlir/Dialect/Rock/IR/GetRockInfo.h

@@ -0,0 +1,61 @@
+//===- GetRockInfo.h - functions used to calculate information about Rock ops
+//---------------===//
+//
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_ROCK_IR_GETROCKINFO_H
+#define MLIR_DIALECT_ROCK_IR_GETROCKINFO_H
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/Rock/IR/Rock.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+
+namespace mlir {
+class Operation;
+class Type;
+
+namespace rock {
+
+// This function returns the func or gpu.func of a given op
+Operation *getParentFuncOp(Operation *op);
+
+// Return a boolean if the features contain accel properties
+bool isAccel(rock::GemmFeatures features);
+
+// Get the arch from the op
+FailureOr<StringAttr> getArch(Operation *op);
+
+// Get the arch from the op and error out if it cannot be found
+StringAttr getArchValue(Operation *op);
+
+// Get the num_cu from the op
+FailureOr<int64_t> getNumCU(Operation *op);
+
+// Get the num_cu from the op, and error out if it cannot be found
+int64_t getNumCUValue(Operation *op);
+
+inline rock::GemmFeatures intersectGemmFeatures(rock::GemmFeatures a,
+                                                rock::GemmFeatures b) {
+  return a & b;
+}
+
+// Get the features enabled for the specified op. These will be dependent on
+// the architecture being used, and the type of the op.
+rock::GemmFeatures getFeatures(Operation *op);
+
+// This function returns a boolean value if the underlying op has support for
+// an optional 'GemmFeatures' attribute
+bool opHasOptionalFeature(Operation *op);
+
+} // End namespace rock
+} // End namespace mlir
+#endif // MLIR_DIALECT_ROCK_IR_GETROCKINFO_H

mlir/include/mlir/Dialect/Rock/IR/RockConvInterface.td

@@ -66,7 +66,7 @@ def RockConvInterface : OpInterface<"RockConvInterface"> {
         /*defaultImplementation=*/[{ return $_op.getDilations(); }]>,
     InterfaceMethod<
         /*desc=*/[{Return features attr values.}],
-        /*retType=*/"GemmFeatures",
+        /*retType=*/"std::optional<GemmFeatures>",
         /*methodName=*/"getFeatures",
         /*args=*/(ins),
         /*methodBody=*/"",

mlir/include/mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.td

@@ -32,20 +32,20 @@ def RockGemmGemmWrapperInterface : OpInterface<"RockGemmGemmWrapperInterface"> {
   let methods = [
     InterfaceMethod<
         /*desc=*/[{
-          Return the KernelType of this op
+          Return the type from this op that is needed to calculate GemmFeatures
         }],
-        /*retType=*/"::mlir::rock::KernelType",
-        /*methodName=*/"getKernelType",
+        /*retType=*/"SmallVector<::mlir::Type>",
+        /*methodName=*/"getTypesForFeature",
         /*args=*/(ins),
-        /*methodBody=*/"",
+        /*methodBody=*/"return {$_op.getAType(), $_op.getCType()};",
         /*defaultImplementation=*/""
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the arch string of this op
+          Return the KernelType of this op
         }],
-        /*retType=*/"StringRef",
-        /*methodName=*/"getArch",
+        /*retType=*/"::mlir::rock::KernelType",
+        /*methodName=*/"getKernelType",
         /*args=*/(ins),
         /*methodBody=*/"",
         /*defaultImplementation=*/""
@@ -154,9 +154,9 @@ def RockGemmGemmWrapperInterface : OpInterface<"RockGemmGemmWrapperInterface"> {
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the features attribute of this op.
+          Return the optional features attribute of this op.
         }],
-        /*retType=*/"::mlir::rock::GemmFeatures",
+        /*retType=*/"std::optional<::mlir::rock::GemmFeatures>",
         /*methodName=*/"getGemmFeatures",
         /*args=*/(ins),
         /*methodBody=*/"",
@@ -166,13 +166,15 @@ def RockGemmGemmWrapperInterface : OpInterface<"RockGemmGemmWrapperInterface"> {
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the optional number of Compute Units the GPU provides.
+          Return the features attribute of this op.
         }],
-        /*retType=*/"std::optional<uint32_t>",
-        /*methodName=*/"getNumCU",
+        /*retType=*/"::mlir::rock::GemmFeaturesAttr",
+        /*methodName=*/"getGemmFeaturesAttr",
         /*args=*/(ins),
         /*methodBody=*/"",
-        /*defaultImplementation=*/ ""
+        /*defaultImplementation=*/[{
+          return $_op.getFeaturesAttr();
+        }]
       >,
 
     InterfaceMethod<

mlir/include/mlir/Dialect/Rock/IR/RockGemmWrapperInterface.td

@@ -30,20 +30,20 @@ def RockGemmWrapperInterface : OpInterface<"RockGemmWrapperInterface"> {
   let methods = [
     InterfaceMethod<
         /*desc=*/[{
-          Return the KernelType of this op
+          Return the type from this op that is needed to calculate GemmFeatures
         }],
-        /*retType=*/"::mlir::rock::KernelType",
-        /*methodName=*/"getKernelType",
+        /*retType=*/"SmallVector<::mlir::Type>",
+        /*methodName=*/"getTypesForFeature",
         /*args=*/(ins),
-        /*methodBody=*/"",
+        /*methodBody=*/"return {$_op.getAType()};",
         /*defaultImplementation=*/""
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the arch string of this op
+          Return the KernelType of this op
         }],
-        /*retType=*/"StringRef",
-        /*methodName=*/"getArch",
+        /*retType=*/"::mlir::rock::KernelType",
+        /*methodName=*/"getKernelType",
         /*args=*/(ins),
         /*methodBody=*/"",
         /*defaultImplementation=*/""
@@ -102,9 +102,9 @@ def RockGemmWrapperInterface : OpInterface<"RockGemmWrapperInterface"> {
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the features attribute of this op.
+          Return the optional features attribute of this op.
         }],
-        /*retType=*/"::mlir::rock::GemmFeatures",
+        /*retType=*/"std::optional<::mlir::rock::GemmFeatures>",
         /*methodName=*/"getGemmFeatures",
         /*args=*/(ins),
         /*methodBody=*/"",
@@ -114,13 +114,15 @@ def RockGemmWrapperInterface : OpInterface<"RockGemmWrapperInterface"> {
       >,
     InterfaceMethod<
         /*desc=*/[{
-          Return the optional number of Compute Units the GPU provides.
+          Return the features attribute of this op.
         }],
-        /*retType=*/"std::optional<uint32_t>",
-        /*methodName=*/"getNumCU",
+        /*retType=*/"::mlir::rock::GemmFeaturesAttr",
+        /*methodName=*/"getGemmFeaturesAttr",
         /*args=*/(ins),
         /*methodBody=*/"",
-        /*defaultImplementation=*/ ""
+        /*defaultImplementation=*/[{
+          return $_op.getFeaturesAttr();
+        }]
       >,
 
     /// Setters for tuning-related data