Fix duplicate mapping detection in gpu::setMappingAttr()

[RoboTux]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-gpu

Author: Thomas Preud'homme (RoboTux)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/77499.diff

4 Files Affected:

• (modified) mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp (+1)
• (modified) mlir/unittests/Dialect/CMakeLists.txt (+1)
• (added) mlir/unittests/Dialect/GPU/CMakeLists.txt (+10)
• (added) mlir/unittests/Dialect/GPU/ParallelLoopMapperTest.cpp (+42)

diff --git a/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp b/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp
index 72e0ebc132e862..9d398998dd63bd 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp
@@ -41,6 +41,7 @@ gpu::setMappingAttr(ParallelOp ploopOp,
         specifiedMappings.count(processor))
       return ploopOp.emitError(
           "invalid mapping multiple loops to same processor");
+    specifiedMappings.insert(processor);
   }
   ArrayRef<Attribute> mappingAsAttrs(mapping.data(), mapping.size());
   ploopOp->setAttr(getMappingAttrName(),
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 2dec4ba3c001e8..48c9b59265e384 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -7,6 +7,7 @@ target_link_libraries(MLIRDialectTests
   MLIRDialect)
 
 add_subdirectory(ArmSME)
+add_subdirectory(GPU)
 add_subdirectory(Index)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
diff --git a/mlir/unittests/Dialect/GPU/CMakeLists.txt b/mlir/unittests/Dialect/GPU/CMakeLists.txt
new file mode 100644
index 00000000000000..2d31cc890918fd
--- /dev/null
+++ b/mlir/unittests/Dialect/GPU/CMakeLists.txt
@@ -0,0 +1,10 @@
+add_mlir_unittest(MLIRGPUTests
+  ParallelLoopMapperTest.cpp
+)
+target_link_libraries(MLIRGPUTests
+  PRIVATE
+  MLIRArithDialect
+  MLIRGPUDialect
+  MLIRGPUTransforms
+  MLIRParser
+  MLIRSCFDialect)
diff --git a/mlir/unittests/Dialect/GPU/ParallelLoopMapperTest.cpp b/mlir/unittests/Dialect/GPU/ParallelLoopMapperTest.cpp
new file mode 100644
index 00000000000000..80b86d53ebcdb3
--- /dev/null
+++ b/mlir/unittests/Dialect/GPU/ParallelLoopMapperTest.cpp
@@ -0,0 +1,42 @@
+//===- ParallelLoopMapper.cpp - Parallel Loop Mapper unit tests -----------===//
+//
+// Part of the LLVM 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Transforms/ParallelLoopMapper.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Parser/Parser.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+TEST(ParallelLoopMapperTest, TestSetMappingAttrMultipleProcMapping) {
+  MLIRContext context;
+  Builder b(&context);
+  context.getOrLoadDialect<arith::ArithDialect>();
+  context.getOrLoadDialect<scf::SCFDialect>();
+  context.getOrLoadDialect<gpu::GPUDialect>();
+
+  const char *const code = R"mlir(
+    %c1 = arith.constant 1 : index
+    %c3 = arith.constant 3 : index
+    scf.parallel (%i) = (%c1) to (%c3) step (%c1) {
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(code, &context);
+  scf::ParallelOp ploopOp =
+      *(*module).getBody()->getOps<scf::ParallelOp>().begin();
+  auto BlockXMapping = b.getAttr<gpu::ParallelLoopDimMappingAttr>(
+      gpu::Processor::BlockX, b.getDimIdentityMap(), b.getDimIdentityMap());
+  SmallVector<gpu::ParallelLoopDimMappingAttr, 1> mapping = {BlockXMapping,
+                                                             BlockXMapping};
+  EXPECT_FALSE(succeeded(gpu::setMappingAttr(ploopOp, mapping)));
+}
+
+} // end anonymous namespace

[joker-eph]

Can this be tested with lit testing?

[RoboTux]

Can this be tested with lit testing?

No, the only caller for that function will never pass a mapping with twice the same processor attribute.

[RoboTux]

Ping?

[joker-eph]

I still have concerns about the testing: this diverges from how we test this kind of things in the project.

[RoboTux]

I still have concerns about the testing: this diverges from how we test this kind of things in the project.

Should I remove the test then? As I said, the only caller to the function being fixed never triggers the wrong codepath so it can only be exercise by adding a new caller (which is what happens when adding the new unit test).

[joker-eph]

Maybe best to skip the unit-tests here.