[Intel] Add permute to TargetInfo

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.cpp

@@ -85,6 +85,14 @@ Value TargetInfo::shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
   return LLVM::intel::shuffleIdx(loc, rewriter, val, i);
 }
 
+Value TargetInfo::permute(RewriterBase &rewriter, Location loc, Value a,
+                          Value b, Value selector) const {
+  // Warning: The `a` and `b` operands are ordered to align with Nvidia's `prmt`
+  // Both use little-endian ordering, but AMD puts the MSBs of the data in the
+  // 0-th operand.
+  return LLVM::intel::permute(loc, rewriter, b, a, selector);
+}
+
 Value TargetInfo::programId(RewriterBase &rewriter, Location loc,
                             ModuleOp moduleOp, ProgramIDDim axis) const {
   Value blockId =

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.h

@@ -41,6 +41,9 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
   Value shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
                    Value i) const override;
 
+  Value permute(RewriterBase &rewriter, Location loc, Value a, Value b,
+                Value selector) const override;
+
   Value programId(RewriterBase &rewriter, Location loc, ModuleOp moduleOp,
                   ProgramIDDim axis) const override;
 

third_party/intel/lib/TritonIntelGPUToLLVM/Utility.cpp

@@ -101,6 +101,23 @@ Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, Value i) {
   return shuffleCommon(loc, rewriter, val, i, mlir::gpu::ShuffleMode::IDX);
 }
 
+Value permute(Location loc, RewriterBase &rewriter, Value x, Value y,
+              Value selector) {
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  Value prmt_mask = selector;
+  // convert from nybble mask to byte mask:
+  prmt_mask =
+      b.or_(b.and_(prmt_mask, b.i32_val(0x000000ff)),
+            b.shl(b.and_(prmt_mask, b.i32_val(0x0000ff00)), b.i32_val(8)));
+  prmt_mask =
+      b.or_(b.and_(prmt_mask, b.i32_val(0x000f000f)),
+            b.shl(b.and_(prmt_mask, b.i32_val(0x00f000f0)), b.i32_val(4)));
+  Value args[] = {x, y, prmt_mask};
+  auto op = createLLVMIntrinsicCallOp(rewriter, loc, "llvm.amdgcn.perm", i32_ty,
+                                      args);
+  return op.getResult(0);
+}
+
 LLVM::RoundingMode
 convertTritonRoundingModeToLLVM(const triton::RoundingMode rounding) {
   LLVM::RoundingMode roundingMode;

third_party/intel/lib/TritonIntelGPUToLLVM/Utility.h

@@ -19,6 +19,9 @@ Value shuffleUp(Location loc, RewriterBase &rewriter, Value val, int i);
 Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, int i);
 Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, Value i);
 
+Value permute(Location loc, RewriterBase &rewriter, Value a, Value b,
+              Value selector);
+
 /// Create a predicated block, using \p cond as the condition and \p ops for the
 /// values supplied by the conditional branch to the exit block. The \p
 /// thenOpsFn function is used to inject operations in the 'then' branch: