[Mlir-commits] [mlir] Fix duplicate mapping detection in gpu::setMappingAttr() (PR #77499)

[Thomas Preud'homme]

https://github.com/RoboTux created https://github.com/llvm/llvm-project/pull/77499

None

>From b912db7ab37a06f9f94f74762de0a58281b270bc Mon Sep 17 00:00:00 2001
From: Thomas Preud'homme <thomas.preudhomme at arm.com>
Date: Tue, 9 Jan 2024 15:49:42 +0000
Subject: [PATCH] Fix duplicate mapping detection in gpu::setMappingAttr()

---
 .../GPU/Transforms/ParallelLoopMapper.cpp     |  1 +
 mlir/unittests/Dialect/CMakeLists.txt         |  1 +
 mlir/unittests/Dialect/GPU/CMakeLists.txt     | 10 +++++
 .../Dialect/GPU/ParallelLoopMapperTest.cpp    | 42 +++++++++++++++++++
 4 files changed, 54 insertions(+)
 create mode 100644 mlir/unittests/Dialect/GPU/CMakeLists.txt
 create mode 100644 mlir/unittests/Dialect/GPU/ParallelLoopMapperTest.cpp

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