Fix incorrect fusion check (#1956)

In AffixTuningParameters we will calculate the GemmFeatures based on the ops in the func, at which point we will then write this as an attribute to the func. The problem with this is that MIGraphX, and other users of rocMLIR, could query isModuleFusible at any point, potentially before AffixTuningParameters has run. This is perfectly legal behavior, but there was a bug in fusionUtils that expected the GemmFeatures to have already been written to the func when isModuleFusible was called.

The fix for this that I came up with was to call rock::getFeatures directly on ReduceOps and teach getFeatures how to handle these ops.

I also added in a new CAPI test as we seemingly did not have one testing fusibility of reduce ops.

Author: justinrosner

mlir/include/mlir/Dialect/Rock/IR/CMakeLists.txt

@@ -20,6 +20,7 @@ mlir_tablegen(RockAccelTuningParamAttrInterface.h.inc -gen-attr-interface-decls)
 mlir_tablegen(RockAccelTuningParamAttrInterface.cpp.inc -gen-attr-interface-defs)
 add_public_tablegen_target(MLIRRockAccelTuningParamAttrInterfaceIncGen)
 
+add_mlir_interface(RockGemmFeaturesInterface)
 add_mlir_interface(RockGemmGemmWrapperInterface)
 add_mlir_interface(RockGemmWrapperInterface)
 add_mlir_interface(RockConvInterface)

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

@@ -52,10 +52,6 @@ inline rock::GemmFeatures intersectGemmFeatures(rock::GemmFeatures a,
 // 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/Rock.h

@@ -82,6 +82,7 @@ constexpr int64_t maxHardwareWorkgroupSize = 1024;
 
 #include "mlir/Dialect/Rock/IR/RockAcceptingViewOpInterface.h"
 #include "mlir/Dialect/Rock/IR/RockConvInterface.h"
+#include "mlir/Dialect/Rock/IR/RockGemmFeaturesInterface.h"
 #include "mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.h"
 #include "mlir/Dialect/Rock/IR/RockGemmWrapperInterface.h"
 #include "mlir/Dialect/Rock/IR/RockWriterOpInterface.h"

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

@@ -10,6 +10,7 @@
 #define ROCK_ATTRS
 
 include "mlir/Dialect/Rock/IR/RockBase.td"
+include "mlir/Dialect/Rock/IR/RockGemmFeaturesInterface.td"
 include "mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.td"
 include "mlir/Dialect/Rock/IR/RockGemmWrapperInterface.td"
 include "mlir/Dialect/Rock/IR/RockTuningParamAttrInterface.td"

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

@@ -0,0 +1,25 @@
+//===---- RockGemmFeaturesInterface.h - ops that wrap rock.gemm -*- C++ -*-===//
+//
+// Part of the rocMLIR 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
+//
+// Copyright (c) 2025 Advanced Micro Devices INc.
+//===----------------------------------------------------------------------===//
+//
+// This file defines RockGemmFeaturesInterface, which abstracts Rock ops for
+// which we would expect to extract GemmFeatures for.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ROCK_IR_ROCKGEMMFEATURESINTERFACE_H
+#define MLIR_DIALECT_ROCK_IR_ROCKGEMMFEATURESINTERFACE_H
+
+#include "mlir/Dialect/Rock/IR/GemmSize.h"
+#include "mlir/IR/OpDefinition.h"
+
+#include "mlir/Dialect/Rock/IR/RockTypes.h"
+
+#include "mlir/Dialect/Rock/IR/RockGemmFeaturesInterface.h.inc"
+
+#endif // MLIR_DIALECT_ROCK_IR_ROCKGEMMFEATURESINTERFACE_H

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

@@ -0,0 +1,38 @@
+//===------------------ RockGemmFeaturesInterface.td ----------------------===//
+//
+// Part of the rocMLIR 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
+//
+// Copyright (c) 2025 Advanced Micro Devices Inc.
+//===----------------------------------------------------------------------===//
+//
+// This file defines RockGemmFeaturesInterface, which abstracts Rock ops for
+// which we would expect to extract GemmFeatures for.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ROCK_GEMM_FEATURES_INTERFACE
+#define ROCK_GEMM_FEATURES_INTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def RockGemmFeaturesInterface : OpInterface<"RockGemmFeaturesInterface"> {
+  let description = [{
+    Interface to abstract away operations for which we need to extract
+    information relating to GemmFeatures from.
+  }];
+  let cppNamespace = "::mlir::rock";
+
+  let methods = [InterfaceMethod<
+      /*desc=*/[{
+          Return the type from this op that is needed to calculate GemmFeatures
+        }],
+      /*retType=*/"SmallVector<::mlir::Type>",
+      /*methodName=*/"getTypesForFeature",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"">];
+}
+
+#endif // ROCK_GEMM_FEATURES_INTERFACE

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

@@ -30,16 +30,6 @@ def RockGemmGemmWrapperInterface : OpInterface<"RockGemmGemmWrapperInterface"> {
   let cppNamespace = "::mlir::rock";
 
   let methods = [
-    InterfaceMethod<
-        /*desc=*/[{
-          Return the type from this op that is needed to calculate GemmFeatures
-        }],
-        /*retType=*/"SmallVector<::mlir::Type>",
-        /*methodName=*/"getTypesForFeature",
-        /*args=*/(ins),
-        /*methodBody=*/"return {$_op.getAType(), $_op.getCType()};",
-        /*defaultImplementation=*/""
-      >,
     InterfaceMethod<
         /*desc=*/[{
           Return the KernelType of this op

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

@@ -28,16 +28,6 @@ def RockGemmWrapperInterface : OpInterface<"RockGemmWrapperInterface"> {
   let cppNamespace = "::mlir::rock";
 
   let methods = [
-    InterfaceMethod<
-        /*desc=*/[{
-          Return the type from this op that is needed to calculate GemmFeatures
-        }],
-        /*retType=*/"SmallVector<::mlir::Type>",
-        /*methodName=*/"getTypesForFeature",
-        /*args=*/(ins),
-        /*methodBody=*/"return {$_op.getAType()};",
-        /*defaultImplementation=*/""
-      >,
     InterfaceMethod<
         /*desc=*/[{
           Return the KernelType of this op

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

@@ -15,6 +15,7 @@
 
 include "mlir/Dialect/Rock/IR/RockAttrDefs.td"
 include "mlir/Dialect/Rock/IR/RockConvInterface.td"
+include "mlir/Dialect/Rock/IR/RockGemmFeaturesInterface.td"
 include "mlir/Dialect/Rock/IR/RockGemmGemmWrapperInterface.td"
 include "mlir/Dialect/Rock/IR/RockGemmWrapperInterface.td"
 include "mlir/Dialect/Rock/IR/RockAcceptingViewOpInterface.td"
@@ -61,11 +62,13 @@ class TensorOrMemRefOf<list<Type> allowedTypes> :
 
 class IndexArrayLength<int n> : ConfinedAttr<IndexArrayAttr, [ArrayMinCount<n>]>;
 
-class Rock_ConvOpBase<string mnemonic, list<Type> inputTypes=[F32, F16, BF16], list<Type> outputTypes=[F32, F16, BF16]> :
-    Rock_Op<mnemonic, [DeclareOpInterfaceMethods<RockGemmWrapperInterface>,
-                       DeclareOpInterfaceMethods<RockConvInterface>,
-                       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-                       RockFusionRoot]>{
+class Rock_ConvOpBase<string mnemonic, list<Type> inputTypes = [F32, F16, BF16],
+                      list<Type> outputTypes = [F32, F16, BF16]>
+    : Rock_Op<mnemonic, [DeclareOpInterfaceMethods<RockGemmWrapperInterface>,
+                         DeclareOpInterfaceMethods<RockConvInterface>,
+                         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                         DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
+                         RockFusionRoot]> {
   dag commonConvArgs = (ins TensorOrMemRefOf<inputTypes>:$filter,
       TensorOrMemRefOf<inputTypes>:$input,
       TensorOrMemRefOf<outputTypes>:$output,
@@ -132,6 +135,7 @@ def Rock_ConvBwdWeightOp : Rock_ConvOpBase<"conv_bwd_weight">
 
 def Rock_GemmOp
     : Rock_Op<"gemm", [DeclareOpInterfaceMethods<RockGemmWrapperInterface>,
+                       DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
                        RockFusionRoot]>,
       Arguments<(
           ins Arg<TensorOrMemRefOf<GemmInputTypes>, "matrix A", [MemRead]>:$a,
@@ -187,12 +191,17 @@ def Rock_ReduceOp
   }];
   let extraClassDeclaration = [{
     ::mlir::OpOperand* getOutArgument() { return &(*this)->getOpOperand(1); }
+
+    SmallVector<::mlir::Type> getTypesForFeature() {
+      return {getIn().getType()};
+    }
   }];
 }
 def Rock_AttentionOp
     : Rock_Op<
           "attention", [DeclareOpInterfaceMethods<RockGemmGemmWrapperInterface>,
                         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                        DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
                         RockFusionRoot, AttrSizedOperandSegments,
                         AttrSizedResultSegments]>,
       Arguments<(ins TensorOrMemRefOf<[F32, F16, BF16, I8]>:$queries,
@@ -260,6 +269,7 @@ def Rock_GemmElementwiseGemmOp
     : Rock_Op<"gemm_elementwise_gemm",
               [DeclareOpInterfaceMethods<RockGemmGemmWrapperInterface>,
                DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+               DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
                RockFusionRoot]>,
       AllElementTypesMatch<["a", "b"]>,
       Arguments<(ins TensorOrMemRefOf<[F32, F16, BF16]>:$a,
@@ -310,6 +320,7 @@ def Rock_ConvElementwiseGemmOp
     : Rock_Op<"conv_elementwise_gemm",
               [DeclareOpInterfaceMethods<RockGemmGemmWrapperInterface>,
                DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+               DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
                RockFusionRoot]>,
       AllElementTypesMatch<["filter", "input"]>,
       Arguments<(ins TensorOrMemRefOf<[F32, F16, BF16]>:$filter,
@@ -486,7 +497,9 @@ def Rock_TensorUntransformCastOp :
 }
 
 def Rock_GridwiseGemmOp
-    : Rock_Op<"gridwise_gemm", [RockFusionRoot]>,
+    : Rock_Op<"gridwise_gemm", [DeclareOpInterfaceMethods<
+                                    RockGemmFeaturesInterface>,
+                                RockFusionRoot]>,
       Arguments<(ins Arg<MemRefRankOf<GemmInputTypes, [3]>,
                          "matrix A view", [MemRead]>:$a,
           Arg<MemRefRankOf<GemmInputTypes, [3]>, "matrix B view", [MemRead]>:$b,
@@ -503,19 +516,13 @@ def Rock_GridwiseGemmOp
     $c `=` $a `*` $b `storeMethod` `(` $storeMethod `)` (`features` `=` $features^)? attr-dict `:` type($c) `=` type($a) `*` type($b)
   }];
   let hasVerifier = 1;
-
-  // Return the type from this op that is needed to calculate GemmFeatures
-  let extraClassDeclaration = [{
-    SmallVector<::mlir::Type> getTypesForFeature() {
-      SmallVector<::mlir::Type> types = {getA().getType()};
-      return types;
-    }
-  }];
 }
 
 // gridwise_gemm_accel
 def Rock_GridwiseGemmAccelOp
-    : Rock_Op<"gridwise_gemm_accel", [RockFusionRoot]>,
+    : Rock_Op<"gridwise_gemm_accel", [DeclareOpInterfaceMethods<
+                                          RockGemmFeaturesInterface>,
+                                      RockFusionRoot]>,
       Arguments<(ins Arg<MemRefRankOf<GemmInputTypes, [3]>,
                          "matrix A view", [MemRead]>:$a,
           Arg<MemRefRankOf<GemmInputTypes, [3]>, "matrix B view", [MemRead]>:$b,
@@ -532,21 +539,15 @@ def Rock_GridwiseGemmAccelOp
     `(` operands `)` `storeMethod` `(` $storeMethod `)` (`features` `=` $features^)? attr-dict `:` type(operands)
   }];
   let hasVerifier = 1;
-
-  // Return the type from this op that is needed to calculate GemmFeatures
-  let extraClassDeclaration = [{
-    SmallVector<::mlir::Type> getTypesForFeature() {
-      SmallVector<::mlir::Type> types = {getA().getType()};
-      return types;
-    }
-  }];
 }
 
 // gridwise_attention_accel
 def Rock_GridwiseAttentionAccelOp
     : Rock_Op<"gridwise_attention_accel",
               [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-               RockFusionRoot, AttrSizedOperandSegments]>,
+               DeclareOpInterfaceMethods<RockGemmFeaturesInterface>,
+               RockFusionRoot, AttrSizedOperandSegments,
+]>,
       Arguments<(ins MemRefRankOf<[F32, F16, BF16, I8], [3]>:$queries,
           MemRefRankOf<[F32, F16, BF16, I8], [3]>:$keys,
           MemRefRankOf<[F32, F16, BF16], [3]>:$values,
@@ -572,15 +573,6 @@ def Rock_GridwiseAttentionAccelOp
     `(` operands `)` (`features` `=` $features^)? `preSoftmaxOps` `=` $preSoftmaxBody attr-dict `:` type(operands)
   }];
   let hasVerifier = 1;
-
-  // Return the type from this op that is needed to calculate GemmFeatures
-  let extraClassDeclaration = [{
-    SmallVector<::mlir::Type> getTypesForFeature() {
-      SmallVector<::mlir::Type> types = {getKeys().getType(),
-                                         getValues().getType()};
-      return types;
-    }
-  }];
 }
 
 // Memory allocation on GPU memory hierachy.
@@ -1370,7 +1362,8 @@ defvar AccelResTypes = [VectorOfLengthAndType<[4, 8, 16, 32], [F32, I32, F16, BF
 
 // blockwise_gemm_accel
 def Rock_BlockwiseGemmAccelOp
-    : Rock_Op<"blockwise_gemm_accel">,
+    : Rock_Op<"blockwise_gemm_accel", [DeclareOpInterfaceMethods<
+                                          RockGemmFeaturesInterface>]>,
       Arguments<(ins MemRefOf<LdsBufferTypes>:$matrixA,
           MemRefOf<LdsBufferTypes>:$matrixB, I32Attr:$inMPerThread,
           I32Attr:$inNPerThread, UnitAttr:$rotateMWithK, UnitAttr:$rotateNWithK,
@@ -1393,14 +1386,6 @@ def Rock_BlockwiseGemmAccelOp
     `:` type($matrixC) `+` `` `=` type($bufferA) `from` type($matrixA) `*`
                                   type($bufferB) `from` type($matrixB)
   }];
-
-  // Return the type from this op that is needed to calculate GemmFeatures
-  let extraClassDeclaration = [{
-    SmallVector<::mlir::Type> getTypesForFeature() {
-      SmallVector<::mlir::Type> types = {getMatrixA().getType()};
-      return types;
-    }
-  }];
 }
 
 // threadwise_gemm
@@ -1426,7 +1411,8 @@ def Rock_ThreadwiseGemmOp:
 }
 // threadwise_accel_gemm
 def Rock_ThreadwiseAccelGemmOp
-    : Rock_Op<"threadwise_accel_gemm">,
+    : Rock_Op<"threadwise_accel_gemm", [DeclareOpInterfaceMethods<
+                                           RockGemmFeaturesInterface>]>,
       Arguments<(ins Arg<MemRefOf<NativeMemoryOpTypes>,
                          "source register view A", [MemRead]>:$matrixA,
           Arg<MemRefOf<NativeMemoryOpTypes>,
@@ -1448,14 +1434,6 @@ def Rock_ThreadwiseAccelGemmOp
     `:` type($matrixC) `+` `` `=` type($matrixA) `*` type($matrixB)
   }];
   let hasVerifier = 1;
-
-  // Return the type from this op that is needed to calculate GemmFeatures
-  let extraClassDeclaration = [{
-    SmallVector<::mlir::Type> getTypesForFeature() {
-      SmallVector<::mlir::Type> types = {getMatrixA().getType()};
-      return types;
-    }
-  }];
 }
 
 // blockwise_broadcasting_reduction

mlir/lib/Dialect/Rock/IR/CMakeLists.txt

@@ -4,6 +4,7 @@ add_rocmlir_dialect_library(MLIRRockOps
   GetRockInfo.cpp
   TransformMapBuilder.cpp
   RockDialect.cpp
+  RockGemmFeaturesInterface.cpp
   RockGemmWrapperInterface.cpp
   RockGemmGemmWrapperInterface.cpp
   RockConvInterface.cpp
@@ -19,6 +20,7 @@ add_rocmlir_dialect_library(MLIRRockOps
 
   DEPENDS
   MLIRRockAttrDefsIncGen
+  MLIRRockGemmFeaturesInterfaceIncGen
   MLIRRockGemmWrapperInterfaceIncGen
   MLIRRockGemmGemmWrapperInterfaceIncGen
   MLIRRockTuningParamAttrInterfaceIncGen