Address review

Author: nicolasvasilache

mlir/include/mlir/Dialect/GPU/IR/GPUDeviceMappingAttr.td

@@ -257,16 +257,9 @@ def GPUMappingMaskAttr : GPU_Attr<"GPUMappingMask", "mask", [
   let parameters = (ins "uint64_t":$mask);
   let assemblyFormat = "`<` params `>`";
   let description = [{
-    Attribute describing how to filter the processing units that a
-    region is mapped to.
-
-    In the first implementation the masking is a bitfield that specifies for
-    each processing unit whether it is active or not.
-
-    In the future, we may want to implement this as a symbol to refer to
-    dynamically defined values.
-
-    Extending op semantics with an operand is deemed too intrusive at this time.
+    Attribute describing how to filter the processing units that a region is
+    mapped to. The masking is a bitfield that specifies for each processing
+    unit whether it is active or not.
   }];
 }
 

mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h

@@ -11,7 +11,6 @@
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
-#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
@@ -57,7 +56,7 @@ mapForallToBlocksImpl(RewriterBase &rewriter, TransformOpInterface transformOp,
 DiagnosedSilenceableFailure
 mapOneForallToThreadsImpl(RewriterBase &rewriter,
                           std::optional<TransformOpInterface> transformOp,
-                          scf::ForallOp forallOp, ArrayRef<int64_t> blockDims,
+                          scf::ForallOp forallOp, ArrayRef<int64_t> blockSizes,
                           int64_t warpSize, bool syncAfterDistribute);
 
 /// Search `scf.forall` ops nested under `target` and map each such op to an

mlir/include/mlir/Dialect/GPU/TransformOps/Utils.h

@@ -77,6 +77,7 @@ struct GpuIdBuilder {
 /// used for indexing rewrites as well as 3D sizes for predicate generation.
 /// If `useLinearMapping` is true, the `idBuilder` method returns nD values
 /// used for indexing rewrites as well as 1D sizes for predicate generation.
+/// If `mask` is provided, it will be used to filter the active blocks.
 struct GpuBlockIdBuilder : public GpuIdBuilder {
   GpuBlockIdBuilder(MLIRContext *ctx, bool useLinearMapping = false,
                     DeviceMaskingAttrInterface mask = nullptr);
@@ -87,6 +88,7 @@ struct GpuBlockIdBuilder : public GpuIdBuilder {
 /// used for indexing rewrites as well as 3D sizes for predicate generation.
 /// If `useLinearMapping` is true, the `idBuilder` method returns nD values
 /// used for indexing rewrites as well as 1D sizes for predicate generation.
+/// If `mask` is provided, it will be used to filter the active warpgroups.
 struct GpuWarpgroupIdBuilder : public GpuIdBuilder {
   GpuWarpgroupIdBuilder(MLIRContext *ctx, int64_t warpSize,
                         bool useLinearMapping = false,
@@ -101,6 +103,7 @@ struct GpuWarpgroupIdBuilder : public GpuIdBuilder {
 /// used for indexing rewrites as well as 3D sizes for predicate generation.
 /// If `useLinearMapping` is true, the `idBuilder` method returns nD values
 /// used for indexing rewrites as well as 1D sizes for predicate generation.
+/// If `mask` is provided, it will be used to filter the active warps.
 struct GpuWarpIdBuilder : public GpuIdBuilder {
   GpuWarpIdBuilder(MLIRContext *ctx, int64_t warpSize,
                    bool useLinearMapping = false,
@@ -113,6 +116,7 @@ struct GpuWarpIdBuilder : public GpuIdBuilder {
 /// used for indexing rewrites as well as 3D sizes for predicate generation.
 /// If `useLinearMapping` is true, the `idBuilder` method returns nD values
 /// used for indexing rewrites as well as 1D sizes for predicate generation.
+/// If `mask` is provided, it will be used to filter the active threads.
 struct GpuThreadIdBuilder : public GpuIdBuilder {
   GpuThreadIdBuilder(MLIRContext *ctx, bool useLinearMapping = false,
                      DeviceMaskingAttrInterface mask = nullptr);
@@ -122,6 +126,7 @@ struct GpuThreadIdBuilder : public GpuIdBuilder {
 /// The `idBuilder` method returns nD values used for indexing rewrites as well
 /// as 1D sizes for predicate generation.
 /// This `useLinearMapping` case is the only supported case.
+/// If `mask` is provided, it will be used to filter the active lanes.
 struct GpuLaneIdBuilder : public GpuIdBuilder {
   GpuLaneIdBuilder(MLIRContext *ctx, int64_t warpSize, bool unused,
                    DeviceMaskingAttrInterface mask = nullptr);

mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td

@@ -66,18 +66,35 @@ def DeviceMaskingAttrInterface : AttrInterface<"DeviceMaskingAttrInterface"> {
     Attribute interface describing how to filter the processing units that a
     region is mapped to.
 
-    A popcount can be applied to determine the logical linear index that a
-    physical processing unit is responsible for.
+    For instance, consider the following example mask which specifies processing
+    units 2, 4 and 5 are active:
+    ```
+              8       4       0
+      mask  : 0 0 0 1 1 0 1 0 0
+    ```
+    The logical ID for an active processing unit is defined as its position
+    relative to the other active processing units. In this example, we have:
+    ```
+      Processing Unit    LogicalID
+              0             N/A
+              1             N/A
+              2              0
+              3             N/A
+              4              1
+              5              2
+              6             N/A
+              7             N/A
+    ```
   }];
 
  let methods = [
     InterfaceMethod<
       /*desc=*/[{
-        Return the logical active id for a given physical id.
+        Create the logical active id for a given physical id.
         Expects a physicalLinearMappingId of I64Type.
       }],
       /*retTy=*/"Value",
-      /*methodName=*/"getLogicalLinearMappingId",
+      /*methodName=*/"createLogicalLinearMappingId",
       /*args=*/(ins "OpBuilder&":$builder, "Value":$physicalLinearMappingId)
     >,
     InterfaceMethod<
@@ -87,7 +104,7 @@ def DeviceMaskingAttrInterface : AttrInterface<"DeviceMaskingAttrInterface"> {
         Expects a physicalLinearMappingId of I64Type.
       }],
       /*retTy=*/"Value",
-      /*methodName=*/"getIsActiveIdPredicate",
+      /*methodName=*/"createIsActiveIdPredicate",
       /*args=*/(ins "OpBuilder&":$builder, "Value":$physicalLinearMappingId)
     >,
     InterfaceMethod<

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -133,7 +133,7 @@ int64_t GPUMappingMaskAttr::getMaxNumPhysicalIds() const { return 64; }
 /// Example filter: 0 0 0 0 1 1 1 1 1
 /// Intersection  : 0 0 0 0 1 0 1 0 0
 /// PopCnt        : 2
-Value GPUMappingMaskAttr::getLogicalLinearMappingId(
+Value GPUMappingMaskAttr::createLogicalLinearMappingId(
     OpBuilder &b, Value physicalLinearMappingId) const {
   Location loc = physicalLinearMappingId.getLoc();
   Value mask = b.create<arith::ConstantOp>(loc, b.getI64IntegerAttr(getMask()));
@@ -154,7 +154,7 @@ Value GPUMappingMaskAttr::getLogicalLinearMappingId(
 /// Example filter: 0 0 0 1 0 0 0 0 0
 /// Intersection  : 0 0 0 1 0 0 0 0 0
 /// Cmp           : 1
-Value GPUMappingMaskAttr::getIsActiveIdPredicate(
+Value GPUMappingMaskAttr::createIsActiveIdPredicate(
     OpBuilder &b, Value physicalLinearMappingId) const {
   Location loc = physicalLinearMappingId.getLoc();
   Value mask = b.create<arith::ConstantOp>(loc, b.getI64IntegerAttr(getMask()));

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

@@ -12,9 +12,7 @@
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/Utils.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
@@ -758,7 +756,8 @@ static DiagnosedSilenceableFailure
 getThreadIdBuilder(std::optional<TransformOpInterface> transformOp,
                    scf::ForallOp forallOp, ArrayRef<int64_t> blockSizes,
                    int64_t warpSize, GpuIdBuilder &gpuIdBuilder) {
-  auto mappingAttr = forallOp.getDeviceMappingAttrs().front();
+  DeviceMappingAttrInterface mappingAttr =
+      forallOp.getDeviceMappingAttrs().front();
   bool useLinearMapping = mappingAttr.isLinearMapping();
 
   // Sanity checks that may result in runtime verification errors.

mlir/lib/Dialect/GPU/TransformOps/Utils.cpp

@@ -160,7 +160,7 @@ commonLinearIdBuilderFn(int64_t multiplicity = 1,
       scaledLinearIdI64 = rewriter.create<arith::IndexCastUIOp>(
           loc, rewriter.getI64Type(), scaledLinearId);
       Value logicalLinearIdI64 =
-          mask.getLogicalLinearMappingId(rewriter, scaledLinearIdI64);
+          mask.createLogicalLinearMappingId(rewriter, scaledLinearIdI64);
       scaledLinearId = rewriter.create<arith::IndexCastUIOp>(
           loc, rewriter.getIndexType(), logicalLinearIdI64);
       LDBG("------adjusting linearId with mask: " << scaledLinearId);
@@ -184,7 +184,7 @@ commonLinearIdBuilderFn(int64_t multiplicity = 1,
     // 4. If mask present, it takes precedence to determine predication.
     if (mask) {
       Value isActiveIdPredicate =
-          mask.getIsActiveIdPredicate(rewriter, scaledLinearIdI64);
+          mask.createIsActiveIdPredicate(rewriter, scaledLinearIdI64);
       LDBG("------adjusting predicate with mask: " << isActiveIdPredicate);
       predicateOps.push_back(isActiveIdPredicate);
     } else {