[mlir][Transform] Add support for mapping to linearized GPU ids

[nicolasvasilache]

Summary

Depends on: D146130

[nicolasvasilache]

Update

[nicolasvasilache]

Update tests.

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

let description = [{
     ```mlir
     %laneId = gpu.lane_id
     ```
   }];
   let results = (outs Index:$result);
   let assemblyFormat = "attr-dict";
 }
 
+def GPU_LinearIdOp : GPU_Op<"linear_id", [
+      Pure, DeclareOpInterfaceMethods<InferIntRangeInterface>]> {
+  let description = [{
+    Returns the linearized id within the workgroup (block).
+
+    Example:
+    ```mlir
+    %laneId = gpu.lane_id
+    ```
+  }];
+  let results = (outs Index:$result);
+  let assemblyFormat = "attr-dict";
+}
+
 def GPU_SubgroupIdOp : GPU_Op<"subgroup_id", [
       Pure, DeclareOpInterfaceMethods<InferIntRangeInterface>]>,
     Arguments<(ins)>, Results<(outs Index:$result)> {
   let description = [{
     Returns the subgroup id, i.e. the index of the current subgroup within the
     workgroup.
 
     Example:

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

let description = [{
 
     Warp (aka subgroup) are grouped into a grid where grid may be
     described by a 1-, 2-, or 3-dimensional rectangle. This attribute indicates
     that thread block parallelism is desired. It can be consumed by lowering to
     generate GPU code.
   }];
 }
 
+def LinearIdEnum : I64EnumAttr<"LinearId", "threads for loop mapping", [
+    DimX, DimY, DimZ]> {
+  let cppNamespace = "::mlir::gpu";
+}
+
+def GPULinearIdMapping : GPU_Attr<"GPULinearIdMapping", "linear", [
+  DeclareAttrInterfaceMethods<DeviceMappingAttrInterface> ] >  {
+  let parameters = (ins
+    EnumParameter<LinearIdEnum>:$linear_id
+  );
+  let assemblyFormat = "`<` params `>`";
+  let description = [{
+    An attribute that allows defining thread parallelism for GPU devices.
+  }];
+}
+
 def BlocksEnum : I64EnumAttr<"Blocks", "threads for loop mapping", [
     DimX, DimY, DimZ]> {
   let cppNamespace = "::mlir::gpu";
 }
 
 def GPUBlockMappingAttr : GPU_Attr<"GPUBlockMapping", "block", [
   DeclareAttrInterfaceMethods<DeviceMappingAttrInterface> ] >  {
   let parameters = (ins

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

let description = [{
         gpu.terminator
       }
       ```
     }];
 
   let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$block_dims,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$warp_dims,
+                   DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$linear_dims,
                    DefaultValuedAttr<BoolAttr, "true">:$sync_after_distribute);
   let results = (outs PDL_Operation:$result);
 
   let assemblyFormat = [{
     $target
     `block_dims` `=` $block_dims
     (`warp_dims` `=` $warp_dims^)?
+    (`linear_dims` `=` $linear_dims^)?
     (`sync_after_distribute` `=` $sync_after_distribute^)?
     attr-dict
   }];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
-
 def MapForallToBlocks :
   Op<Transform_Dialect, "gpu.map_forall_to_blocks",
     [FunctionalStyleTransformOpTrait,
      MemoryEffectsOpInterface,
      TransformOpInterface,
      TransformEachOpTrait]> {
   let description = [{
     Target the gpu_launch op and rewrite the top level `scf.forall`

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

 int64_t GPUBlockMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getBlock());
 }
 
 int64_t GPUWarpMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getWarp());
 }
 
+int64_t GPULinearIdMappingAttr::getMappingId() const {
+  return static_cast<int64_t>(getLinearId());
+}
+
 int64_t GPUThreadMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getThread());
 }
 
 int64_t GPUMemorySpaceMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getAddressSpace());
 }
 

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

void ThreadIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
   setResultRange(getResult(), getIndexRange(0, max - 1ULL));
 }
 
 void LaneIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
                                  SetIntRangeFn setResultRange) {
   setResultRange(getResult(), getIndexRange(0, kMaxSubgroupSize - 1ULL));
 }
 
+void LinearIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
+                                   SetIntRangeFn setResultRange) {
+  setResultRange(getResult(), getIndexRange(0, kMaxSubgroupSize - 1ULL));
+}
+
 void SubgroupIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
                                      SetIntRangeFn setResultRange) {
   setResultRange(getResult(), getIndexRange(0, kMaxDim - 1ULL));
 }
 
 void GlobalIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
                                    SetIntRangeFn setResultRange) {
   uint64_t blockDimMax =

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

idBuilder = [](RewriterBase &rewriter, scf::ForallOp forallOp,
       LLVM_DEBUG(llvm::interleaveComma(ids, DBGS() << "--ids: ");
                  llvm::dbgs() << "\n");
 
       return ids;
     };
   }
 };
 
+struct GpuLinearIdBuilder : public GpuIdBuilder {
+  GpuLinearIdBuilder(MLIRContext *ctx) : GpuIdBuilder() {
+    mappingAttributes = {GPULinearIdMappingAttr::get(ctx, LinearId::DimX),
+                         GPULinearIdMappingAttr::get(ctx, LinearId::DimY),
+                         GPULinearIdMappingAttr::get(ctx, LinearId::DimZ)};
+    idBuilder = [](RewriterBase &rewriter, scf::ForallOp forallOp,
+                   ArrayRef<int64_t> mappingDims) {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPoint(forallOp);
+      Location loc = forallOp.getLoc();
+      Value linearIdOp = rewriter.create<LinearIdOp>(loc);
+      SmallVector<int64_t> reverseBlockDims(llvm::reverse(mappingDims));
+      LLVM_DEBUG(llvm::interleaveComma(reverseBlockDims,
+                                       DBGS() << "--delinearization basis: ");
+                 llvm::dbgs() << "\n");
+
+      SmallVector<int64_t> strides = computeStrides(reverseBlockDims);
+      LLVM_DEBUG(llvm::interleaveComma(strides,
+                                       DBGS() << "--delinearization strides: ");
+                 llvm::dbgs() << "\n");
+
+      AffineExpr d0;
+      bindDims(rewriter.getContext(), d0);
+      SmallVector<AffineExpr> delinearizingExprs = delinearize(d0, strides);
+      LLVM_DEBUG(llvm::interleaveComma(delinearizingExprs,
+                                       DBGS() << "--delinearization exprs: ");
+                 llvm::dbgs() << "\n");
+
+      SmallVector<Value> ids;
+      for (AffineExpr e : delinearizingExprs)
+        ids.push_back(makeComposedAffineApply(rewriter, loc, e, linearIdOp));
+      LLVM_DEBUG(llvm::interleaveComma(ids, DBGS() << "--ids: ");
+                 llvm::dbgs() << "\n");
+
+      return ids;
+    };
+  }
+};
+
 } // namespace
 
 static DiagnosedSilenceableFailure
 failureHelper(std::optional<TransformOpInterface> transformOp,
               scf::ForallOp forallOp, const Twine &message) {
   if (transformOp.has_value())
     return transformOp->emitSilenceableError() << message;
   return emitDefiniteFailure(forallOp, message);
checkMappingAttributeTypes(std::optional<TransformOpInterface> transformOp,
   bool hasThreadMapping =
       llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
         return attr.isa<GPUThreadMappingAttr>();
       });
   bool hasWarpMapping =
       llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
         return attr.isa<GPUWarpMappingAttr>();
       });
+  bool hasLinearMapping =
+      llvm::any_of(forallOp.getMapping().value(), [](Attribute attr) {
+        return attr.isa<GPULinearIdMappingAttr>();
+      });
   int64_t countMappingTypes = 0;
   countMappingTypes += hasBlockMapping ? 1 : 0;
   countMappingTypes += hasThreadMapping ? 1 : 0;
   countMappingTypes += hasWarpMapping ? 1 : 0;
+  countMappingTypes += hasLinearMapping ? 1 : 0;
   if (countMappingTypes > 1) {
     return failureHelper(transformOp, forallOp,
                          "cannot mix different mapping types, use nesting");
   }
 
   DenseSet<Attribute> seen;
   for (Attribute map : forallOp.getMapping()->getValue()) {
     if (llvm::is_contained(seen, map)) {
DiagnosedSilenceableFailure transform::MapNestedForallToThreads::applyToOne(
     Operation *target, ApplyToEachResultList &results, TransformState &state) {
   LaunchOp gpuLaunch = dyn_cast<LaunchOp>(target);
   auto transformOp = cast<TransformOpInterface>(getOperation());
 
   // Basic high-level verifications.
   if (!gpuLaunch)
     return emitSilenceableError() << "Given target is not a gpu.launch";
 
+  // Mapping to block ids.
   SmallVector<int64_t> blockDims{getBlockDims()};
   if (blockDims.size() != 3)
     return transformOp.emitDefiniteFailure("requires size-3 thread mapping");
   DiagnosedSilenceableFailure diag =
       checkGpuLimits(transformOp, std::nullopt, std::nullopt, std::nullopt,
                      blockDims[0], blockDims[1], blockDims[2]);
   if (diag.isSilenceableFailure()) {
     diag.attachNote(getLoc()) << getBlockDimsAttrName() << " is too large";
LLVM_DEBUG(
       llvm::dbgs() << "\n");
   GpuThreadIdBuilder gpuThreadIdBuilder(ctx);
   diag = mlir::transform::gpu::mapNestedForallToThreadsImpl(
       rewriter, transformOp, target, blockDims, getSyncAfterDistribute(),
       gpuThreadIdBuilder);
   if (!diag.succeeded())
     return diag;
 
+  // Mapping to warp ids.
   SmallVector<int64_t> warpDims{getWarpDims()};
   if (!warpDims.empty()) {
     if (warpDims.size() != 3)
       return transformOp.emitDefiniteFailure("requires size-3 warp mapping");
   }
 
   LLVM_DEBUG(llvm::interleaveComma(
                  warpDims, DBGS() << "mapNestedForallToThreadsImpl warpDims: ");
target->walk([&](SubgroupIdOp subgroupIdOp) {
     AffineExpr ltid = getAffineDimExpr(0, ctx);
     auto warpId = makeComposedFoldedAffineApply(rewriter, subgroupIdOp.getLoc(),
                                                 ltid.floorDiv(kWarpSize),
                                                 {linearThreadId});
     LDBG("----warpId: " << warpId);
     rewriter.replaceAllUsesWith(subgroupIdOp, warpId.get<Value>());
   });
 
+  // Mapping to linear ids.
+  SmallVector<int64_t> linearDims{getLinearDims()};
+  if (!linearDims.empty()) {
+    if (linearDims.size() != 3)
+      return transformOp.emitDefiniteFailure("requires size-3 linear mapping");
+  }
+
+  LLVM_DEBUG(
+      llvm::interleaveComma(
+          linearDims, DBGS() << "mapNestedForallToThreadsImpl linearDims: ");
+      llvm::dbgs() << "\n");
+  GpuLinearIdBuilder gpuLinearIdBuilder(ctx);
+  diag = mlir::transform::gpu::mapNestedForallToThreadsImpl(
+      rewriter, transformOp, target, linearDims, getSyncAfterDistribute(),
+      gpuLinearIdBuilder);
+  if (!diag.succeeded())
+    return diag;
+
+  // Perform a late replacement of the LinearIdOp, taking into account the
+  // blockDims.
+  target->walk([&](LinearIdOp linearIdOp) {
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(linearIdOp);
+    auto linearThreadId =
+        getStaticLinearThreadId(rewriter, linearIdOp.getLoc(), blockDimsOfr);
+    LDBG("----linearThreadId: " << linearThreadId);
+    rewriter.replaceAllUsesWith(linearIdOp, linearThreadId.get<Value>());
+  });
+
   results.push_back(gpuLaunch.getOperation());
   return diag;
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/GPU/transform-gpu.mlir

^bb1(%arg0: !pdl.operation):
   transform.gpu.map_nested_forall_to_threads %funcop block_dims = [12, 9, 1] sync_after_distribute = false
 }
 
 // -----
 
 !type = memref<2 x 32 x f32>
 !type1d = memref<32 x f32>
 
-// CHECK-DAG: #[[MAPWY:.*]] = affine_map<(d0, d1, d2) -> (((d0 + d1 * 12 + d2 * 132) floordiv 32) floordiv 4)>
-// CHECK-DAG: #[[MAPWX:.*]] = affine_map<(d0, d1, d2) -> ((((d0 + d1 * 12 + d2 * 132) floordiv 32) mod 4) floordiv 2)>
+// CHECK-DAG: #[[$MAPWY:.*]] = affine_map<(d0, d1, d2) -> (((d0 + d1 * 12 + d2 * 132) floordiv 32) floordiv 4)>
+// CHECK-DAG: #[[$MAPWX:.*]] = affine_map<(d0, d1, d2) -> ((((d0 + d1 * 12 + d2 * 132) floordiv 32) mod 4) floordiv 2)>
+
+// CHECK-DAG: #[[$MAPLX:.*]] = affine_map<(d0, d1, d2) -> (d2 * 4 + (d0 + d1 * 12) floordiv 33)>
+// CHECK-DAG: #[[$MAPLY:.*]] = affine_map<(d0, d1) -> (((d0 + d1 * 12) mod 33) floordiv 11)>
+// CHECK-DAG: #[[$MAPLZ:.*]] = affine_map<(d0, d1) -> ((d0 + d1 * 12) mod 11)>
 
 // CHECK-LABEL: func.func @map_multi_level(
 func.func @map_multi_level(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {
   %one = arith.constant 1 : index
-  %c12 = arith.constant 12 : index
+  %c10 = arith.constant 10 : index
   %c9 = arith.constant 9 : index
   %c7 = arith.constant 7 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
 
-  // CHECK: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
 
   // check that both the thread level and the warp level got distributed.
   //  CHECK-NOT: #gpu.thread
   //  CHECK-NOT: #gpu.warp
   %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)
             threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)
   {
     // CHECK-DAG: %[[TIDX:.*]] = gpu.thread_id  x
     // CHECK-DAG: %[[TIDY:.*]] = gpu.thread_id  y
     scf.forall (%i, %j) in (%c7, %c9) {
       %4 = memref.load %x[%i, %j] : !type
       %5 = memref.load %y[%i, %j] : !type
       %6 = math.fma %alpha, %4, %5 : f32
       memref.store %6, %y[%i, %j] : !type
     }  { mapping = [#gpu.thread<y>, #gpu.thread<x>]}
 
     // CHECK-DAG: %[[TIDZ:.*]] = gpu.thread_id  z
-    // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[MAPWY]](%[[TIDX]], %[[TIDY]], %[[TIDZ]])
-    // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[MAPWX]](%[[TIDX]], %[[TIDY]], %[[TIDZ]])
+    // CHECK-DAG: %[[WIDY:.*]] = affine.apply #[[$MAPWY]](%[[TIDX]], %[[TIDY]], %[[TIDZ]])
+    // CHECK-DAG: %[[WIDX:.*]] = affine.apply #[[$MAPWX]](%[[TIDX]], %[[TIDY]], %[[TIDZ]])
     // CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[WIDX]], %[[C1]] : index
     // CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[WIDY]], %[[C1]] : index
     //     CHECK: %[[COND:.*]] = arith.andi %[[CMPY]], %[[CMPX]] : i1
     //     CHECK: scf.if %[[COND]]
     scf.forall (%i) in (%c1) {
         %7 = memref.load %t[%i] : !type1d
         %8 = arith.addf %alpha, %7 : f32
         memref.store %8, %t[%i] : !type1d
      }  {mapping = [#gpu.warp<x>] }
+
+    // CHECK-DAG: %[[LIDX:.*]] = affine.apply #[[$MAPLX]](%[[TIDX]], %[[TIDY]], %[[TIDZ]])
+    // CHECK-DAG: %[[LIDY:.*]] = affine.apply #[[$MAPLY]](%[[TIDX]], %[[TIDY]])
+    // CHECK-DAG: %[[LIDZ:.*]] = affine.apply #[[$MAPLZ]](%[[TIDX]], %[[TIDY]])
+    // CHECK-DAG: %[[CMPX:.*]] = arith.cmpi ult, %[[LIDX]], %[[C10]] : index
+    // CHECK-DAG: %[[CMPY:.*]] = arith.cmpi ult, %[[LIDY]], %[[C2]] : index
+    //     CHECK: %[[CONDXY:.*]] = arith.andi %[[CMPX]], %[[CMPY]] : i1
+    // CHECK-DAG: %[[CMPZ:.*]] = arith.cmpi ult, %[[LIDZ]], %[[C1]] : index
+    //     CHECK: %[[CONDXYZ:.*]] = arith.andi %[[CONDXY]], %[[CMPZ]] : i1
+    //     CHECK: scf.if %[[CONDXYZ]]
+    scf.forall (%i, %j) in (%c10, %c2) {
+        %7 = memref.load %t[%i] : !type1d
+        %8 = arith.addf %alpha, %7 : f32
+        memref.store %8, %t[%j] : !type1d
+     }  {mapping = [#gpu.linear<x>, #gpu.linear<y>] }
     gpu.terminator
   }
   return %y : !type
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!pdl.operation) -> !pdl.operation
   transform.gpu.map_nested_forall_to_threads %funcop
-    block_dims = [12, 11, 1] warp_dims = [2, 2, 1]
+    block_dims = [12, 11, 1] warp_dims = [2, 2, 1] linear_dims = [11, 3, 4]
 }