[Codegen][Tuner] Expose XOR shuffle bounds and validation functions in CAPI (#23442)

After #23175, we now generate
swizzle hint ops for scaled gemms whose parameters are set during
`lowering_config` selection. These parameters (`rowElems` and
`accessElems`) can be chosen via the tuner too (although in the future
we intend for there to be an analytically derived solution to this).
This PR exposes two functions, `getXorShuffleBounds` and
`isXorShuffleValid`, to allow the tuner to constrain its search space
for applicable XOR swizzles.

Assisted by: composer-1

---------

Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>

Author: Muzammiluddin-Syed-ECE

compiler/bindings/c/iree/compiler/dialects/iree_gpu.h

@@ -184,6 +184,19 @@ ireeGPUTargetInfoGetMMAIntrinsics(MlirAttribute mmaIntrinsics,
                                   mma_intrinsic_enum_t *mmaIntrinsicVals,
                                   uint8_t *virtualMmaIntrinsicTags);
 
+// Returns the lower and upper bounds for valid XOR shuffle parameters for a
+// given MMA intrinsic and operand index. On success, writes to minAccessElems
+// and totalTileElems and returns true. Returns false on failure.
+MLIR_CAPI_EXPORTED bool ireeGPUGetXorShuffleBounds(MlirAttribute mmaIntrinsic,
+                                                   int32_t operandIndex,
+                                                   int64_t *minAccessElems,
+                                                   int64_t *totalTileElems);
+
+// Returns true if the XOR shuffle is valid for the given parameters.
+MLIR_CAPI_EXPORTED bool ireeGPUIsXORShuffleValid(int64_t numRowElems,
+                                                 int64_t numAccessElems,
+                                                 int64_t totalTileElems);
+
 #ifdef __cplusplus
 }
 #endif

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -485,7 +485,7 @@ NB_MODULE(_ireeCompilerDialects, m) {
             if (!mlirAttributeIsNull(basisInfo.mappingAttr)) {
               mapping = getIntArrayAttrValues(basisInfo.mappingAttr);
             }
-            return std::make_tuple(counts, mapping);
+            return std::tuple(counts, mapping);
           })
       .def_property_readonly(
           "mma_kind", [](MlirAttribute self) -> std::optional<MlirAttribute> {
@@ -633,14 +633,38 @@ NB_MODULE(_ireeCompilerDialects, m) {
       });
 
   iree_gpu_module.def(
-      "get_single_subgroup_layout",
-      [](MlirAttribute attr, int fragment) {
-        return ireeGPUGetSingleSubgroupLayout(attr, fragment);
-      },
+      "get_single_subgroup_layout", ireeGPUGetSingleSubgroupLayout,
       "Returns the single subgroup layout (element, thread, outer, "
       "tstrides) for a given MMA or VirtualMMA intrinsic and fragment. ",
       py::arg("attr"), py::arg("fragment"));
 
+  //===-------------------------------------------------------------------===//
+  // Binding to XOR shuffle utility functions
+  //===-------------------------------------------------------------------===//
+
+  iree_gpu_module.def(
+      "get_xor_shuffle_bounds",
+      [](MlirAttribute mmaIntrinsic,
+         int operandIndex) -> std::optional<std::tuple<int64_t, int64_t>> {
+        int64_t minAccessElems = 0, totalTileElems = 0;
+        if (ireeGPUGetXorShuffleBounds(mmaIntrinsic, operandIndex,
+                                       &minAccessElems, &totalTileElems)) {
+          return std::tuple(minAccessElems, totalTileElems);
+        }
+        return std::nullopt;
+      },
+      "Returns the bounds for valid XOR shuffle parameters (min_access_elems, "
+      "total_tile_elems) for the given MMA intrinsic and operand index. See "
+      "GPUUtils for sweep semantics. Returns (min_access_elems, "
+      "total_tile_elems) or None on failure.",
+      py::arg("mmaIntrinsic"), py::arg("operand_index"));
+
+  iree_gpu_module.def(
+      "is_xor_shuffle_valid", ireeGPUIsXORShuffleValid,
+      "Returns true if the XOR shuffle is valid for the given parameters.",
+      py::arg("num_row_elems"), py::arg("num_access_elems"),
+      py::arg("total_tile_elems"));
+
   //===-------------------------------------------------------------------===//
   // Binding to utility function getExecutableVariantOps
   //===-------------------------------------------------------------------===//

compiler/bindings/python/test/api/tuner_api_test.py

@@ -6,6 +6,7 @@
 
 from iree.compiler import ir
 from iree.compiler.dialects import iree_codegen
+from iree.compiler.dialects import iree_gpu
 from iree.compiler.dialects import affine
 from iree.compiler.ir import AffineMap, AffineDimExpr
 
@@ -545,3 +546,34 @@ def test_infer_scaled_contraction_dimensions():
     assert dims_batched.n == [2], f"Got {dims_batched.n}"
     assert dims_batched.k == [3], f"Got {dims_batched.k}"
     assert dims_batched.kB == [4], f"Got {dims_batched.kB}"
+
+
+@run
+def test_is_xor_shuffle_valid():
+    """Test XOR shuffle validation (pure function, no MLIR attributes)."""
+    # Valid: row and access divide tile; row >= access; tile >= row.
+    assert iree_gpu.is_xor_shuffle_valid(256, 32, 512)
+    assert iree_gpu.is_xor_shuffle_valid(512, 64, 512)
+    assert iree_gpu.is_xor_shuffle_valid(32, 8, 512)
+    # Invalid: row exceeds tile.
+    assert not iree_gpu.is_xor_shuffle_valid(512, 32, 256)
+    # Invalid: access exceeds row.
+    assert not iree_gpu.is_xor_shuffle_valid(256, 512, 512)
+    # Invalid: row does not evenly divide tile.
+    assert not iree_gpu.is_xor_shuffle_valid(300, 32, 512)
+    # Invalid: access does not evenly divide row.
+    assert not iree_gpu.is_xor_shuffle_valid(256, 33, 512)
+
+
+@run
+def test_get_xor_shuffle_bounds():
+    """Test XOR shuffle bounds for an MMA intrinsic (for use by SharkTuner)."""
+    # Use an MMA intrinsic that supports getXorShuffleBounds (InnerTileDescAttrInterface).
+    mma_attr = iree_gpu.MMAAttr.get(iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16)
+    bounds = iree_gpu.get_xor_shuffle_bounds(mma_attr, operand_index=0)
+    assert bounds is not None, "get_xor_shuffle_bounds should succeed for MMAAttr"
+    min_access_elems, total_tile_elems = bounds
+    assert min_access_elems == 4
+    assert total_tile_elems == 256
+    bounds_rhs = iree_gpu.get_xor_shuffle_bounds(mma_attr, operand_index=1)
+    assert bounds_rhs is not None

compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp

@@ -530,3 +530,29 @@ void ireeGPUTargetInfoGetMMAIntrinsics(MlirAttribute mmaIntrinsics,
     assert(false && "Unexpected attribute type in MMA intrinsics array");
   }
 }
+
+bool ireeGPUGetXorShuffleBounds(MlirAttribute mmaIntrinsic,
+                                int32_t operandIndex, int64_t *minAccessElems,
+                                int64_t *totalTileElems) {
+  assert(!mlirAttributeIsNull(mmaIntrinsic) && "mmaIntrinsic cannot be null");
+  auto innerTileDesc = llvm::dyn_cast<
+      mlir::iree_compiler::IREE::Codegen::InnerTileDescAttrInterface>(
+      unwrap(mmaIntrinsic));
+  assert(innerTileDesc && "innerTileDesc cannot be null");
+  assert(minAccessElems && "minAccessElems cannot be null");
+  assert(totalTileElems && "totalTileElems cannot be null");
+  mlir::FailureOr<mlir::iree_compiler::XorShuffleBounds> result =
+      mlir::iree_compiler::getXorShuffleBounds(innerTileDesc, operandIndex);
+  if (llvm::failed(result)) {
+    return false;
+  }
+  *minAccessElems = result->minAccessElems;
+  *totalTileElems = result->totalTileElems;
+  return true;
+}
+
+bool ireeGPUIsXORShuffleValid(int64_t numRowElems, int64_t numAccessElems,
+                              int64_t totalTileElems) {
+  return mlir::iree_compiler::isXORShuffleValid(numRowElems, numAccessElems,
+                                                totalTileElems);
+}

compiler/src/iree/compiler/API/api_exports.c

@@ -28,6 +28,8 @@ extern void ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID();
 extern void ireeCodegenDispatchLoweringPassPipelineAttrGetValue();
 extern void ireeCodegenGetExecutableVariantOps();
 extern void ireeGPUGetSingleSubgroupLayout();
+extern void ireeGPUGetXorShuffleBounds();
+extern void ireeGPUIsXORShuffleValid();
 extern void ireeCodegenGetTunerRootOps();
 extern void ireeCodegenGetAttentionOpDetail();
 extern void ireeCodegenInferScaledContractionDimensions();
@@ -949,6 +951,8 @@ uintptr_t __iree_compiler_hidden_force_extern() {
   x += (uintptr_t)&ireeCodegenDispatchLoweringPassPipelineAttrGetValue;
   x += (uintptr_t)&ireeCodegenGetExecutableVariantOps;
   x += (uintptr_t)&ireeGPUGetSingleSubgroupLayout;
+  x += (uintptr_t)&ireeGPUGetXorShuffleBounds;
+  x += (uintptr_t)&ireeGPUIsXORShuffleValid;
   x += (uintptr_t)&ireeCodegenGetTunerRootOps;
   x += (uintptr_t)&ireeCodegenGetAttentionOpDetail;
   x += (uintptr_t)&ireeCodegenInferScaledContractionDimensions;

compiler/src/iree/compiler/API/api_exports.def

@@ -18,6 +18,8 @@ EXPORTS
   ireeCodegenDispatchLoweringPassPipelineAttrGetValue
   ireeCodegenGetExecutableVariantOps
   ireeGPUGetSingleSubgroupLayout
+  ireeGPUGetXorShuffleBounds
+  ireeGPUIsXORShuffleValid
   ireeCodegenGetTunerRootOps
   ireeCodegenGetAttentionOpDetail
   ireeCodegenInferScaledContractionDimensions

compiler/src/iree/compiler/API/api_exports.ld

@@ -19,6 +19,8 @@ VER_0 {
     ireeCodegenDispatchLoweringPassPipelineAttrGetValue;
     ireeCodegenGetExecutableVariantOps;
     ireeGPUGetSingleSubgroupLayout;
+    ireeGPUGetXorShuffleBounds;
+    ireeGPUIsXORShuffleValid;
     ireeCodegenGetTunerRootOps;
     ireeCodegenGetAttentionOpDetail;
     ireeCodegenInferScaledContractionDimensions;

compiler/src/iree/compiler/API/api_exports.macos.lst

@@ -17,6 +17,8 @@ _ireeCodegenDispatchLoweringPassPipelineAttrGetTypeID
 _ireeCodegenDispatchLoweringPassPipelineAttrGetValue
 _ireeCodegenGetExecutableVariantOps
 _ireeGPUGetSingleSubgroupLayout
+_ireeGPUGetXorShuffleBounds
+_ireeGPUIsXORShuffleValid
 _ireeCodegenGetTunerRootOps
 _ireeCodegenGetAttentionOpDetail
 _ireeCodegenInferScaledContractionDimensions

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

@@ -840,7 +840,7 @@ validateXorShuffle(FailureOr<XorShuffleParams> swizzle,
 // Disabling clang-tidy for the following functions, as it will be externally
 // linked to the CAPI in a future PR.
 // NOLINTBEGIN(misc-use-internal-linkage)
-FailureOr<XorShuffleParams>
+FailureOr<XorShuffleBounds>
 getXorShuffleBounds(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                     int operandIndex) {
   FailureOr<int64_t> maybeMinimumAccessElems =
@@ -850,8 +850,7 @@ getXorShuffleBounds(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
   if (failed(maybeMinimumAccessElems) || failed(maybeTotalTileElems)) {
     return failure();
   }
-  return XorShuffleParams({/*rowElems=*/*maybeMinimumAccessElems,
-                           /*accessElems=*/*maybeTotalTileElems});
+  return XorShuffleBounds{*maybeMinimumAccessElems, *maybeTotalTileElems};
 }
 
 bool isXORShuffleValid(int64_t numRowElems, int64_t numAccessElems,

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h

@@ -211,6 +211,13 @@ struct XorShuffleParams {
   int64_t accessElems;
 };
 
+/// Bounds for valid XOR shuffle parameters (min access elements per thread,
+/// total elements in the tile). Used when sweeping over valid configs.
+struct XorShuffleBounds {
+  int64_t minAccessElems;
+  int64_t totalTileElems;
+};
+
 /// For a given MMA intrinsic and operand, returns the lower bound and upper
 /// bound for valid values of XOR shuffle attribute parameters, access width and
 /// row width. For both parameters, the elements ingested per thread at a time
@@ -221,7 +228,7 @@ struct XorShuffleParams {
 /// the upper bound.
 /// - sweep row elements over all multiple of the access elements, respecting
 /// the upper bound.
-FailureOr<XorShuffleParams>
+FailureOr<XorShuffleBounds>
 getXorShuffleBounds(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                     int operandIndex);