[Mlir-commits] [mlir] [mlir] Graceful handling of non-multiple bit widths for AMDGPU swizzle bitmode lowering (PR #183580)

[Arjun Bhamra]

https://github.com/abhamra updated https://github.com/llvm/llvm-project/pull/183580

>From cf05446b1b00809b41daf8d25cb13c5e2a98dd8b Mon Sep 17 00:00:00 2001
From: Arjun Bhamra <arjun.bhamra25 at gmail.com>
Date: Thu, 26 Feb 2026 12:34:09 -0500
Subject: [PATCH] added fix for swizzle bitwidth compatibility, test

---
 .../Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp | 18 ++++++++++++++++++
 .../Dialect/AMDGPU/conversion_invalid.mlir     |  8 ++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 mlir/test/Dialect/AMDGPU/conversion_invalid.mlir

diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index e7dd8ea6f9149..d148d9bb4137f 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -2522,6 +2522,24 @@ struct AMDGPUSwizzleBitModeLowering
     Location loc = op.getLoc();
     Type i32 = rewriter.getI32Type();
     Value src = adaptor.getSrc();
+
+    // Check early if src's bit width is divisible by 32
+    // for bit swizzling, otherwise it is unable to be decomposed.
+    Type srcType = src.getType();
+    unsigned srcBitWidth;
+    // src can be int, float, or 1D vector
+    if (auto vecTy = llvm::dyn_cast<VectorType>(srcType)) {
+      srcBitWidth = vecTy.getNumElements() *
+                    vecTy.getElementType().getIntOrFloatBitWidth();
+    } else {
+      srcBitWidth = srcType.getIntOrFloatBitWidth();
+    }
+
+    if (srcBitWidth % 32 != 0) {
+      return rewriter.notifyMatchFailure(
+          op, "swizzle_bitmode requires src bit width to be a multiple of 32");
+    }
+
     SmallVector<Value> decomposed =
         LLVM::decomposeValue(rewriter, loc, src, i32);
     unsigned andMask = op.getAndMask();
diff --git a/mlir/test/Dialect/AMDGPU/conversion_invalid.mlir b/mlir/test/Dialect/AMDGPU/conversion_invalid.mlir
new file mode 100644
index 0000000000000..746eda6ff7a13
--- /dev/null
+++ b/mlir/test/Dialect/AMDGPU/conversion_invalid.mlir
@@ -0,0 +1,8 @@
+// RUN: mlir-opt %s -split-input-file --convert-amdgpu-to-rocdl 2>&1 | FileCheck %s
+
+// CHECK: failed to legalize operation 'amdgpu.swizzle_bitmode'
+func.func @swizzle_bitmode_non_multiple_of_32() {
+  %5 = vector.constant_mask [42] : vector<42xi1>
+  %6 = amdgpu.swizzle_bitmode %5 1 2 4 : vector<42xi1>
+  return
+}