fix

Author: dhernandez0

mlir/lib/Dialect/Rock/Transforms/PreparePipeline.cpp

@@ -264,12 +264,11 @@ static bool heuristicPipelineOuterLoop(scf::ForOp outerLoop, ArrayRef<scf::ForOp
   }
 
   // We only pipeline the outer loop for schedule v2
-  if(outerLoop->hasAttr(PipelineAttr::getMnemonic())) {
-    int64_t ii =
-            cast<rock::PipelineAttr>(outerLoop->getAttr(PipelineAttr::getMnemonic()))
-                .getInitiationInterval();
-    pipelineOuterLoop &= ii == 1;
-  }
+  assert(outerLoop->hasAttr(PipelineAttr::getMnemonic()));
+  int64_t ii =
+          cast<rock::PipelineAttr>(outerLoop->getAttr(PipelineAttr::getMnemonic()))
+              .getInitiationInterval();
+  pipelineOuterLoop &= ii == 1;
 
   return pipelineOuterLoop;
 }
@@ -336,10 +335,6 @@ static LogicalResult prepareGemmGemmPipeline(scf::ForOp outerLoop, ArrayRef<scf:
     LLVM_DEBUG(llvm::dbgs()
                 << "Preparing to pipeline inner loops\n");
 
-    // Check if this pass has been already done (applicability)
-    if(!outerLoop->hasAttr(PipelineAttr::getMnemonic()))
-      return success();
-
     // 1. Remove outer loop stages
     SmallVector<rock::StageOp> stages;
     outerLoop.walk([&](rock::StageOp op) {

mlir/tools/rocmlir-tuning-driver/rocmlir-tuning-driver.cpp

@@ -472,33 +472,34 @@ static LogicalResult runTuningLoop(ModuleOp source) {
       rock::createTunableParamSpace(source, tuningSpaceKind));
   for (rock::RockTuningParamAttrInterface tuningAttr :
        tuningSpace->tuningRange) {
-    OwningOpRef<ModuleOp> tuneCopy = cast<ModuleOp>(source->clone());
+    OwningOpRef<ModuleOp> applicabilityCopy = cast<ModuleOp>(source->clone());
     // TODO: remove this once perf_config gets parsed earlier
     SmallString<64> perfConfig;
     tuningAttr.getPerfConfigStr(perfConfig);
     llvm::outs() << perfConfig << "\t";
     StringAttr perfConfigAttr = StringAttr::get(ctx, perfConfig);
-    tuneCopy->walk([&perfConfigAttr](rock::RockGemmWrapperInterface op) {
+    applicabilityCopy->walk([&perfConfigAttr](rock::RockGemmWrapperInterface op) {
       op->setAttr("perf_config", perfConfigAttr);
     });
-    tuneCopy->walk([&perfConfigAttr](rock::RockGemmGemmWrapperInterface op) {
+    applicabilityCopy->walk([&perfConfigAttr](rock::RockGemmGemmWrapperInterface op) {
       op->setAttr("perf_config", perfConfigAttr);
     });
+    OwningOpRef<ModuleOp> compileCopy = cast<ModuleOp>(applicabilityCopy->clone());
 
-    if (!rock::isModuleFusible(tuneCopy.get(), perfConfig)) {
+    if (!rock::isModuleFusible(applicabilityCopy.get(), perfConfig)) {
       llvm::outs() << "N/A\n";
       continue;
     }
 
-    if (failed(applicability.run(tuneCopy.get()))) {
+    if (failed(applicability.run(applicabilityCopy.get()))) {
       llvm::outs() << "N/A\n";
       continue;
     }
 
     SmallVector<uint32_t> blockSizes;
     SmallVector<uint32_t> gridSizes;
     for (auto &fnName : kernelFuncNames) {
-      auto tunedFunc = tuneCopy->lookupSymbol<func::FuncOp>(fnName);
+      auto tunedFunc = applicabilityCopy->lookupSymbol<func::FuncOp>(fnName);
       if (!tunedFunc) {
         llvm::errs() << "Tuned copy somehow missing kernel function\n";
         return failure();
@@ -511,7 +512,7 @@ static LogicalResult runTuningLoop(ModuleOp source) {
     // We have to get these now, they disappear later. Also, if these attributes
     // aren't set the contract of the applicability pipeline changed and that's
     // a problem.
-    if (failed(compilation.run(tuneCopy.get()))) {
+    if (failed(compilation.run(compileCopy.get()))) {
       llvm::errs() << "Backend pipeline failed for config: " << perfConfig
                    << "\n";
       return failure();
@@ -520,7 +521,7 @@ static LogicalResult runTuningLoop(ModuleOp source) {
     // Extract binary and benchmark
     SmallVector<std::string> hipModules;
     for (const auto &fnName : kernelFuncNames) {
-      auto binary = tuneCopy->lookupSymbol<gpu::BinaryOp>(fnName + "_module");
+      auto binary = compileCopy->lookupSymbol<gpu::BinaryOp>(fnName + "_module");
       if (!binary) {
         llvm::errs() << "could not find the GPU binary\n";
       }