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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arjun Bhamra (abhamra)

<details>
<summary>Changes</summary>

This PR handles the case where users call `--convert-amdgpu-to-rocdl` without correct bitwidths. Originally, when the bitwidth was not a multiple of 32, this would fall through to `LLVM:decomposeValue` and trigger an assertion error that `srcBitWidth % dstBitWidth != 0`. To combat this, we check early in the `SwizzleBitModeOp`'s lowering whether the source's bit width is divisible by 32 and fail the pass conditionally.

Closes #<!-- -->180461

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


2 Files Affected:

- (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+18) 
- (added) mlir/test/Dialect/AMDGPU/conversion_invalid.mlir (+8) 


``````````diff
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 3c2c61b2426e9..c1a5b1278f523 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -2530,6 +2530,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
+}

``````````

</details>


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