[mlir][GPU] Integer range inference for GPU index ops

[krzysz00]

Summary

All GPU indices (thread ID, number of threads in a block, etc.) are
semantically non-negative, so we implement InferIntRangeInterface to
communicate this fact.

Depends on D124022

Depends on D124023

Depends on D124024

[Mogball]

mlir/include/mlir/Dialect/GPU/InferIntRangeInterfaceImpl.h

17–18

mlir/include/mlir/InitAllDialects.h

109–111

mlir/test/Dialect/GPU/infer-int-range.mlir

25

could you add smaller tests for all the other ops you added the interface to?

[krzysz00]

Abondoning, CUDA/HIP/... don't give a <= signed_max guarantee on workgroup ID. I will, however, be implementing InferIntRangeINterface downstream for our wrapper around the workgroup ID, where we _can_ make such guarantees (by way of saying "kernel's too big, sorry")

mlir/include/mlir/Dialect/GPU/InferIntRangeInterfaceImpl.h

+//===- InferIntRangeInterfaceImpl.h - Impl. of InferIntRangeInterface ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_GPU_INFERINTRANGEINTERFACEIMPL_H
+#define MLIR_DIALECT_GPU_INFERINTRANGEINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace gpu {
+void registerInferIntRangeInterfaceExternalModels(DialectRegistry &registry);
+} // namespace gpu
+} // namespace mlir
+
+#endif // MLIR_DIALECT_GPU_INFERINTRANGEINTERFACE_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/GPU/InferIntRangeInterfaceImpl.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
registry.insert<acc::OpenACCDialect,
                   x86vector::X86VectorDialect>();
   // clang-format on
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
+
+  gpu::registerInferIntRangeInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUOps
   MLIRMemRef
   MLIRSideEffectInterfaces
   MLIRSupport
   )
 
 add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
+  Transforms/InferIntRangeInterfaceImpl.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
   Transforms/SerializeToBlob.cpp
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/lib/Dialect/GPU/Transforms/InferIntRangeInterfaceImpl.cpp

+//===- InferIntRangeInterfaceImpl.cpp - Impl. of InferIntRangeInterface ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/InferIntRangeInterfaceImpl.h"
+
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+using namespace mlir::arith;
+using namespace mlir::gpu;
+
+namespace {
+/// Declare that index parameters such as block and thread IDs should be
+/// interpereted as unsigned, but do not impose a bound, such as them being
+/// below size_max on them.
+template <typename Op>
+struct IndexOpUnsigned
+    : InferIntRangeInterface::ExternalModel<IndexOpUnsigned<Op>, Op> {
+  using InferIntRangeInterface::ExternalModel<IndexOpUnsigned<Op>,
+                                              Op>::ExternalModel;
+  void inferResultRanges(Operation *rawOp, ArrayRef<IntRangeAttrs> argRanges,
+                         SmallVectorImpl<IntRangeAttrs> &resultRanges) const {
+    Op theOp = cast<Op>(rawOp);
+    resultRanges.push_back(
+        {IntegerAttr::get(theOp->getResultTypes()[0], 0), {}});
+  }
+};
+} // end namespace
+
+void mlir::gpu::registerInferIntRangeInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, GPUDialect *dialect) {
+    BlockDimOp::attachInterface<IndexOpUnsigned<BlockDimOp>>(*ctx);
+    BlockIdOp::attachInterface<IndexOpUnsigned<BlockIdOp>>(*ctx);
+    ThreadIdOp::attachInterface<IndexOpUnsigned<ThreadIdOp>>(*ctx);
+    LaneIdOp::attachInterface<IndexOpUnsigned<LaneIdOp>>(*ctx);
+    GlobalIdOp::attachInterface<IndexOpUnsigned<GlobalIdOp>>(*ctx);
+    SubgroupIdOp::attachInterface<IndexOpUnsigned<SubgroupIdOp>>(*ctx);
+    NumSubgroupsOp::attachInterface<IndexOpUnsigned<NumSubgroupsOp>>(*ctx);
+    SubgroupSizeOp::attachInterface<IndexOpUnsigned<SubgroupSizeOp>>(*ctx);
+  });
+}

mlir/test/Dialect/GPU/infer-int-range.mlir

+// RUN: mlir-opt -arith-fold-inferred-constants -canonicalize -arith-unsigned-when-equivalent %s | FileCheck %s
+
+// Generated from the function body
+//     %1 = affine.apply affine_map<(d0) -> (d0 floordiv 9 + ((d0 mod 9) floordiv 3) * 3 + (d0 mod 3) * 9)>(%0)
+//     return %1 : index
+// via -lower-affine -canonicalize
+// CHECK-LABEL: func @like_affine_map
+// CHECK-DAG: %[[c3:.*]] = arith.constant 3
+// CHECK-DAG: %[[c9:.*]] = arith.constant 9
+// CHECK-DAG: %[[blk:.*]] = gpu.block_id
+// CHECK: %[[i0:.*]] = arith.divsi %[[blk]], %[[c9]]
+// CHECK: %[[i1a:.*]] = arith.remsi %[[blk]], %[[c9]]
+// CHECK: %[[i1b:.*]] = arith.divui %[[i1a]], %[[c3]]
+// CHECK: %[[i1:.*]] = arith.muli %[[i1b]], %[[c3]]
+// CHECK: %[[part1:.*]] = arith.addi %[[i0]], %[[i1]]
+// CHECK: %[[i2a:.*]] = arith.remsi %[[blk]], %[[c3]]
+// CHECK: %[[i2:.*]] = arith.muli %[[i2a]], %[[c9]]
+// CHECK: %[[ret:.*]] = arith.addi %[[part1]], %[[i2]]
+// CHECK: return %[[ret]]
+func @like_affine_map() -> index {
+    %c9 = arith.constant 9 : index
+    %c0 = arith.constant 0 : index
+    %c3 = arith.constant 3 : index
+    %c-1 = arith.constant -1 : index
+    %0 = gpu.block_id  x
+    %1 = arith.cmpi slt, %0, %c0 : index
+    %2 = arith.subi %c-1, %0 : index
+    %3 = arith.select %1, %2, %0 : index
+    %4 = arith.divsi %3, %c9 : index
+    %5 = arith.subi %c-1, %4 : index
+    %6 = arith.select %1, %5, %4 : index
+    %7 = arith.remsi %0, %c9 : index
+    %8 = arith.cmpi slt, %7, %c0 : index
+    %9 = arith.addi %7, %c9 : index
+    %10 = arith.select %8, %9, %7 : index
+    %11 = arith.cmpi slt, %10, %c0 : index
+    %12 = arith.subi %c-1, %10 : index                                                      %13 = arith.select %11, %12, %10 : index
+    %14 = arith.divsi %13, %c3 : index
+    %15 = arith.subi %c-1, %14 : index
+    %16 = arith.select %11, %15, %14 : index
+    %17 = arith.muli %16, %c3 : index
+    %18 = arith.addi %6, %17 : index
+    %19 = arith.remsi %0, %c3 : index
+    %20 = arith.cmpi slt, %19, %c0 : index
+    %21 = arith.addi %19, %c3 : index
+    %22 = arith.select %20, %21, %19 : index
+    %23 = arith.muli %22, %c9 : index
+    %24 = arith.addi %18, %23 : index
+    return %24 : index
+}