[water] add wave.permute op

Signed-off-by: Tim Gymnich <[email protected]>

Author: tgymnich

water/include/water/Dialect/Wave/IR/WaveInterfaces.h

@@ -146,7 +146,7 @@ llvm::LogicalResult verifyElementTypesMatch(std::optional<mlir::Location> loc,
 // location is provided.
 llvm::LogicalResult verifyTypesCompatible(
     wave::WaveTensorType lhs, wave::WaveTensorType rhs,
-    bool includeAddressSpace,
+    bool includeAddressSpace, bool includeShape = true,
     std::optional<mlir::Location> errorLocation = std::nullopt,
     llvm::StringRef lhsName = "", llvm::StringRef rhsName = "");
 
@@ -161,11 +161,10 @@ verifyTypesMatchingDimensions(std::optional<mlir::Location> loc,
                               llvm::ArrayRef<int> rhsDims);
 
 // Verification logic for the compatible-operands traits. Succeeds if all wave
-// tensor-typed operands and results have compatible shapes and, if the
-// corresponding flag is set, compatible address spaces.
-llvm::LogicalResult
-verifyCompatibleOperandsAndResultsOpTrait(mlir::Operation *op,
-                                          bool includeAddressSpace);
+// tensor-typed operands and results have compatible shapes (unless includeShape
+// is false) and, if the corresponding flag is set, compatible address spaces.
+llvm::LogicalResult verifyCompatibleOperandsAndResultsOpTrait(
+    mlir::Operation *op, bool includeAddressSpace, bool includeShape = true);
 }; // namespace detail
 
 template <typename OpTy>
@@ -190,6 +189,17 @@ class CompatibleOperandsAndResultsIgnoreSpaceOpTrait
   }
 };
 
+template <typename OpTy>
+class CompatibleOperandsAndResultsIgnoreShapeOpTrait
+    : public mlir::OpTrait::TraitBase<
+          OpTy, CompatibleOperandsAndResultsIgnoreShapeOpTrait> {
+public:
+  static llvm::LogicalResult verifyTrait(mlir::Operation *op) {
+    return detail::verifyCompatibleOperandsAndResultsOpTrait(
+        op, /*includeAddressSpace=*/true, /*includeShape=*/false);
+  }
+};
+
 //-----------------------------------------------------------------------------
 // WaveElementsPerThreadOpInterface
 //-----------------------------------------------------------------------------

water/include/water/Dialect/Wave/IR/WaveInterfaces.td

@@ -138,6 +138,11 @@ def CompatibleOperandsAndResultsIgnoreSpaceOpTrait
   let cppNamespace = "::wave";
 }
 
+def CompatibleOperandsAndResultsIgnoreShapeOpTrait
+    : NativeOpTrait<"CompatibleOperandsAndResultsIgnoreShapeOpTrait"> {
+  let cppNamespace = "::wave";
+}
+
 //-----------------------------------------------------------------------------
 // WaveInferIndexExprsOpInterface and implementation traits
 //-----------------------------------------------------------------------------

water/include/water/Dialect/Wave/IR/WaveOps.td

@@ -505,4 +505,34 @@ def CastOp : WaveOp<"cast", [
   let hasVerifier = 1;
 }
 
+def PermuteOp : WaveOp<"permute", [
+    DeclareOpInterfaceMethods<WaveInferTypeOpInterface>, CompatibleOperandsAndResultsIgnoreShapeOpTrait,
+    WaveElementsPerThreadOpInterface, IdentityElementsPerThreadOpTrait,
+    DeclareOpInterfaceMethods<WaveInferIndexExprsOpInterface>]> {
+  let summary = "Permute the dimensions of a register-resident tensor";
+  let description = [{
+    Reorders the symbolic dimensions of a register-resident tensor according
+    to the specified target shape. This operation is primarily a semantic
+    marker that affects how index expressions are transformed during
+    compilation. At lowering time, the operation is a pass-through since the
+    actual data layout in registers remains unchanged - only the interpretation
+    of which dimension each element belongs to changes.
+
+    For example, permuting a tensor with shape [B, M, N] to [M, N, B] swaps
+    the strides associated with the symbolic dimensions in the index
+    expressions.
+  }];
+  let arguments = !con((ins
+    Arg<WaveTensorInRegister, "Value to permute">:$value,
+    Arg<WaveSymbolArrayAttr , "Target dimension ordering">:$target_shape
+  ), commonArguments);
+  let results = (outs
+    Res<WaveTensorInRegister, "Permuted value">:$result
+  );
+
+  let assemblyFormat =
+    "$value `,` $target_shape " # commonArgumentsSyntax # " attr-dict `:`"
+    "functional-type(operands, results)";
+}
+
 #endif // WATER_DIALECT_WAVE_WAVEOPS

water/lib/Dialect/Wave/IR/WaveInterfaces.cpp

@@ -375,8 +375,9 @@ llvm::LogicalResult wave::detail::verifyElementTypesMatch(
 
 llvm::LogicalResult wave::detail::verifyTypesCompatible(
     wave::WaveTensorType lhs, wave::WaveTensorType rhs,
-    bool includeAddressSpace, std::optional<Location> errorLocation,
-    llvm::StringRef lhsName, llvm::StringRef rhsName) {
+    bool includeAddressSpace, bool includeShape,
+    std::optional<Location> errorLocation, llvm::StringRef lhsName,
+    llvm::StringRef rhsName) {
   // Fast and cheap path.
   if (lhs == rhs)
     return success();
@@ -402,6 +403,9 @@ llvm::LogicalResult wave::detail::verifyTypesCompatible(
           verifyElementTypesMatch(errorLocation, lhsName, lhs, rhsName, rhs)))
     return failure();
 
+  if (!includeShape)
+    return success();
+
   if (!lhs.getFullySpecified() || !rhs.getFullySpecified())
     return success();
 
@@ -422,7 +426,7 @@ llvm::LogicalResult wave::detail::verifyTypesCompatible(
 static llvm::LogicalResult
 verifyTypeRange(Location loc, TypeRange range,
                 wave::WaveTensorType referenceType, bool includeAddressSpace,
-                llvm::StringRef rangeDescriptionPrefix,
+                bool includeShape, llvm::StringRef rangeDescriptionPrefix,
                 llvm::StringRef referenceDescription) {
   llvm::SmallString<16> rangeDescription(rangeDescriptionPrefix);
   for (auto &&[i, type] : llvm::enumerate(range)) {
@@ -435,16 +439,16 @@ verifyTypeRange(Location loc, TypeRange range,
     os << i;
 
     if (failed(wave::detail::verifyTypesCompatible(
-            tensorType, referenceType, includeAddressSpace, loc, os.str(),
-            referenceDescription))) {
+            tensorType, referenceType, includeAddressSpace, includeShape, loc,
+            os.str(), referenceDescription))) {
       return llvm::failure();
     }
   }
   return llvm::success();
 }
 
 llvm::LogicalResult wave::detail::verifyCompatibleOperandsAndResultsOpTrait(
-    Operation *op, bool includeAddressSpace) {
+    Operation *op, bool includeAddressSpace, bool includeShape) {
   const llvm::StringLiteral kOperandNamePrefix = "operand #";
   const llvm::StringLiteral kResultNamePrefix = "result #";
   std::string referenceDescription;
@@ -470,11 +474,12 @@ llvm::LogicalResult wave::detail::verifyCompatibleOperandsAndResultsOpTrait(
 
   if (llvm::failed(verifyTypeRange(op->getLoc(), op->getOperandTypes(),
                                    referenceType, includeAddressSpace,
-                                   kOperandNamePrefix, os.str())))
+                                   includeShape, kOperandNamePrefix, os.str())))
     return llvm::failure();
 
   return verifyTypeRange(op->getLoc(), op->getResultTypes(), referenceType,
-                         includeAddressSpace, kResultNamePrefix, os.str());
+                         includeAddressSpace, includeShape, kResultNamePrefix,
+                         os.str());
 }
 
 //-----------------------------------------------------------------------------