Merge branch 'main' into use-llvm-func-attrs

Author: victor-eds

python/setup.py

@@ -119,13 +119,13 @@ def find_visual_studio(version_ranges):
     for version_range in version_ranges:
         command = [
             str(vswhere), "-version", version_range, "-requires", "Microsoft.VisualStudio.Component.VC.Tools.x86.x64",
-            "-property", "installationPath", "-prerelease"
+            "-products", "*", "-property", "installationPath", "-prerelease"
         ]
 
         try:
             output = subprocess.check_output(command, text=True).strip()
             if output:
-                return output
+                return output.split("\n")[0]
         except subprocess.CalledProcessError:
             continue
 
@@ -146,6 +146,13 @@ def set_env_vars(vs_path, arch="x64"):
             os.environ[var] = value
 
 
+def initialize_visual_studio_env(version_ranges, arch="x64"):
+    vs_path = find_visual_studio(version_ranges)
+    if not vs_path:
+        raise EnvironmentError("Visual Studio not found in specified version ranges.")
+    set_env_vars(vs_path, arch)
+
+
 # Taken from https://github.com/pytorch/pytorch/blob/master/tools/setup_helpers/env.py
 def check_env_flag(name: str, default: str = "") -> bool:
     return os.getenv(name, default).upper() in ["ON", "1", "YES", "TRUE", "Y"]
@@ -447,10 +454,7 @@ def build_extension(self, ext):
         lit_dir = shutil.which('lit')
         ninja_dir = shutil.which('ninja')
         if platform.system() == "Windows":
-            vs_path = find_visual_studio(["[17.0,18.0)", "[16.0,17.0)"])
-            env = set_env_vars(vs_path)
-            if not vs_path:
-                raise EnvironmentError("Visual Studio 2019 or 2022 not found.")
+            initialize_visual_studio_env(["[17.0,18.0)", "[16.0,17.0)"])
         # lit is used by the test suite
         thirdparty_cmake_args = get_thirdparty_packages([get_llvm_package_info()])
         thirdparty_cmake_args += self.get_pybind11_cmake_args()

python/triton/runtime/CLFinder.py

@@ -19,13 +19,13 @@ def find_visual_studio(version_ranges):
     for version_range in version_ranges:
         command = [
             str(vswhere), "-version", version_range, "-requires", "Microsoft.VisualStudio.Component.VC.Tools.x86.x64",
-            "-property", "installationPath", "-prerelease"
+            "-products", "*", "-property", "installationPath", "-prerelease"
         ]
 
         try:
             output = subprocess.check_output(command, text=True).strip()
             if output:
-                return output
+                return output.split("\n")[0]
         except subprocess.CalledProcessError:
             continue
 
@@ -37,7 +37,7 @@ def set_env_vars(vs_path, arch="x64"):
     if not vcvarsall_path.exists():
         raise FileNotFoundError(f"vcvarsall.bat not found in expected path: {vcvarsall_path}")
 
-    command = f'call "{vcvarsall_path}" {arch} && set'
+    command = ["call", vcvarsall_path, arch, "&&", "set"]
     output = subprocess.check_output(command, shell=True, text=True)
 
     for line in output.splitlines():

test/TritonGEN/tritongen-invalid.mlir

@@ -428,19 +428,3 @@ llvm.func @matrix_2Dblockprefetch(%ptr : !llvm.ptr, %base_width : i32, %base_hei
   triton_gen.2Dblockprefetch %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=32, tile_width=32, tile_height=8, v_blocks=1, cache_control=Default} : (!llvm.ptr, i32, i32, i32, i32, i32)
   llvm.return
 }
-
-// -----
-
-llvm.func @triton_gen.simdblockread(%ptr: !llvm.ptr<3>) {
-  // expected-error @+1 {{'triton_gen.simdblockread' op unsupported vector type}}
-  %ret = triton_gen.simdblockread %ptr : (!llvm.ptr<3>) -> vector<64xi16>
-  llvm.return
-}
-
-// -----
-
-llvm.func @triton_gen.simdblockwrite(%ptr: !llvm.ptr<3>, %val: vector<64xi16>) {
-  // expected-error @+1 {{'triton_gen.simdblockwrite' op unsupported vector type}}
-  triton_gen.simdblockwrite %ptr, %val : (!llvm.ptr<3>, vector<64xi16>)
-  llvm.return
-}

test/TritonGEN/tritongen-to-llvm.mlir

@@ -241,20 +241,29 @@ llvm.func @triton_gen.dpas.bf16_accum(%c: vector<8xbf16>, %a : vector<8xi16>, %b
 
 // CHECK: llvm.func spir_funccc @_Z30intel_sub_group_block_read_us2PU3AS3t(!llvm.ptr<3>) -> vector<2xi16> attributes {memory_effects = #llvm.memory_effects<other = none, argMem = read, inaccessibleMem = none>, no_unwind, will_return}
 
-llvm.func @triton_gen.simdblockread(%ptr: !llvm.ptr<3>) {
-  // CHECK:     llvm.func @triton_gen.simdblockread(%arg0: !llvm.ptr<3>) {
+llvm.func @triton_gen.sub_group_block_read(%ptr: !llvm.ptr<3>) {
+  // CHECK:     llvm.func @triton_gen.sub_group_block_read(%arg0: !llvm.ptr<3>) {
   // CHECK:       llvm.call spir_funccc @_Z30intel_sub_group_block_read_us2PU3AS3t(%arg0) {{.*}} : (!llvm.ptr<3>) -> vector<2xi16>
-  %ret = triton_gen.simdblockread %ptr : (!llvm.ptr<3>) -> vector<2xi16>
+  %ret = triton_gen.sub_group_block_read %ptr : !llvm.ptr<3> -> vector<2xi16>
   llvm.return
 }
 
 // -----
 
 // CHECK: llvm.func spir_funccc @_Z31intel_sub_group_block_write_us2PU3AS3tDv2_t(!llvm.ptr<3>, vector<2xi16>) attributes {memory_effects = #llvm.memory_effects<other = none, argMem = readwrite, inaccessibleMem = none>, no_unwind, will_return}
 
-llvm.func @triton_gen.simdblockwrite(%ptr: !llvm.ptr<3>, %val : vector<2xi16>) {
-  // CHECK:     llvm.func @triton_gen.simdblockwrite(%arg0: !llvm.ptr<3>, %arg1: vector<2xi16>) {
+llvm.func @triton_gen.sub_group_block_write(%ptr: !llvm.ptr<3>, %val : vector<2xi16>) {
+  // CHECK:     llvm.func @triton_gen.sub_group_block_write(%arg0: !llvm.ptr<3>, %arg1: vector<2xi16>) {
   // CHECK:       llvm.call spir_funccc @_Z31intel_sub_group_block_write_us2PU3AS3tDv2_t(%arg0, %arg1) {{.*}} : (!llvm.ptr<3>, vector<2xi16>) -> ()
-  triton_gen.simdblockwrite %ptr, %val : (!llvm.ptr<3>, vector<2xi16>)
+  triton_gen.sub_group_block_write %ptr, %val : !llvm.ptr<3>, vector<2xi16>
+  llvm.return
+}
+
+// -----
+
+llvm.func @triton_gen.sub_group_block_write(%ptr: !llvm.ptr<1>, %val : i32) {
+  // CHECK:     llvm.func @triton_gen.sub_group_block_write(%arg0: !llvm.ptr<1>, %arg1: i32) {
+  // CHECK:       llvm.call spir_funccc @_Z30intel_sub_group_block_write_uiPU3AS1jj(%arg0, %arg1) {{.*}} : (!llvm.ptr<1>, i32) -> ()
+  triton_gen.sub_group_block_write %ptr, %val : !llvm.ptr<1>, i32
   llvm.return
 }

test/TritonGEN/tritongen.mlir

@@ -125,16 +125,16 @@ llvm.func @triton_gen.2Dblockprefetch(%ptr : !llvm.ptr, %base_width : i32, %base
   llvm.return
 }
 
-llvm.func @triton_gen.simdblockread(%ptr : !llvm.ptr) {
-  // CHECK:      llvm.func @triton_gen.simdblockread(%arg0: !llvm.ptr) {
-  // CHECK-NEXT:   triton_gen.simdblockread %arg0 : (!llvm.ptr) -> vector<2xi16>
-  triton_gen.simdblockread %ptr : (!llvm.ptr) -> vector<2xi16>
+llvm.func @triton_gen.sub_group_block_read(%ptr : !llvm.ptr<1>) {
+  // CHECK:      llvm.func @triton_gen.sub_group_block_read(%arg0: !llvm.ptr<1>) {
+  // CHECK-NEXT:   triton_gen.sub_group_block_read %arg0 : !llvm.ptr<1> -> vector<2xi16>
+  triton_gen.sub_group_block_read %ptr : !llvm.ptr<1> -> vector<2xi16>
   llvm.return
 }
 
-llvm.func @triton_gen.simdblockwrite(%ptr : !llvm.ptr, %val : vector<2xi16>) {
-  // CHECK:      llvm.func @triton_gen.simdblockwrite(%arg0: !llvm.ptr, %arg1: vector<2xi16>) {
-  // CHECK-NEXT:    triton_gen.simdblockwrite %arg0, %arg1 : (!llvm.ptr, vector<2xi16>)
-  triton_gen.simdblockwrite %ptr, %val : (!llvm.ptr, vector<2xi16>)
+llvm.func @triton_gen.sub_group_block_write(%ptr : !llvm.ptr<3>, %val : i32) {
+  // CHECK:      llvm.func @triton_gen.sub_group_block_write(%arg0: !llvm.ptr<3>, %arg1: i32) {
+  // CHECK-NEXT:    triton_gen.sub_group_block_write %arg0, %arg1 : !llvm.ptr<3>, i32
+  triton_gen.sub_group_block_write %ptr, %val : !llvm.ptr<3>, i32
   llvm.return
 }

third_party/intel/include/Dialect/TritonGEN/IR/TritonGENOps.td

@@ -314,46 +314,96 @@ def TritonGEN_Matrix2DBlockPrefetchOp : TritonGEN_Op<"2Dblockprefetch">,
   let hasVerifier = 1;
 }
 
-def TritonGEN_SIMDBlockReadOp: TritonGEN_Op<"simdblockread">,
-  Results<(outs FixedVectorOf<[AnyTypeOf<[AnyI8, AnyI16, AnyI32, AnyI64]>]>:$res)>,
-  Arguments<(ins
-    Arg<LLVM_AnyPointer, "", [MemRead]>:$ptr
-  )> {
-
-  let summary = "simd block read";
+def TritonGEN_SubGroupBlockMemoryAccessElementType
+    : AnyTypeOf<[I8, I16, I32, I64],
+                "Valid sub-group block memory access element type">;
+
+def TritonGEN_SubGroupBlockMemoryAccessType
+    : AnyTypeOf<[TritonGEN_SubGroupBlockMemoryAccessElementType,
+                 FixedVectorOfLengthAndType<[2, 4, 8],
+                     [TritonGEN_SubGroupBlockMemoryAccessElementType]>,
+                 // Vectors of length 16 only allowed for i8 for now.
+                 FixedVectorOfLengthAndType<[16], [I8]>],
+          "Valid sub-group block memory access type">;
+
+def TritonGEN_SubGroupBlockMemoryAccessPointerType
+    : Type<And<[LLVM_AnyPointer.predicate,
+                Or<[CPred<"::llvm::cast<::mlir::LLVM::LLVMPointerType>($_self)" #
+                          ".getAddressSpace() == " #
+                          "static_cast<unsigned>(kCrossWorkgroup)">,
+                    CPred<"::llvm::cast<::mlir::LLVM::LLVMPointerType>($_self)" #
+                          ".getAddressSpace() == " #
+                          "static_cast<unsigned>(kWorkgroup)">]>]>,
+          "LLVM pointer in local or global OpenCL address space",
+          "::mlir::LLVM::LLVMPointerType">;
+
+def TritonGEN_SubGroupBlockReadOp: TritonGEN_Op<"sub_group_block_read"> {
+  let summary = "Sub-group block read.";
 
   let description = [{
-    The `triton_gen.simdblockread` operation performs simd block read from
-      a start address without laneId offset. The parameters are:
-      $ptr - the base address to read data
+    The `triton_gen.sub_group_block_read` reads a scalar or vector for each
+    work-item in the sub-group from pointer `ptr` as a block operation.
+    The data is read strided, so the first value is read from:
+    ```
+    ptr[sub_group_local_id]
+    ```
+    and the second one is:
+    ```
+    ptr[sub_group_local_id + sub_group_size]
+    ```
+    etc.
+
+    `ptr` must be aligned to the size of the element type of `res`.
+
+     Example:
+     ```mlir
+     %0 = triton_gen.sub_group_block_read %ptr : !llvm.ptr<1> -> vector<4xi32>
+     ```
   }];
 
+  let arguments = (ins
+      Arg<TritonGEN_SubGroupBlockMemoryAccessPointerType, "", [MemRead]>:$ptr);
+
+  let results = (outs TritonGEN_SubGroupBlockMemoryAccessType:$res);
+
   let assemblyFormat = [{
-    operands ` ` attr-dict `:` functional-type(operands, results)
+    $ptr attr-dict `:` qualified(type($ptr)) `->` type($res)
   }];
-
-  let hasVerifier = 1;
 }
 
-def TritonGEN_SIMDBlockWriteOp : TritonGEN_Op<"simdblockwrite">,
-  Arguments<(ins
-    Arg<LLVM_AnyPointer, "", [MemWrite]>:$ptr,
-    FixedVectorOf<[AnyTypeOf<[AnyI8, AnyI16, AnyI32, AnyI64]>]>:$val
-  )> {
-
+def TritonGEN_SubGroupBlockWriteOp : TritonGEN_Op<"sub_group_block_write"> {
   let summary = "simd block write";
 
   let description = [{
-    The `triton_gen.simdblockwrite` operation performs simd block write to
-      a start address without laneId offset. The parameters are:
-      $ptr - the base address to be written
-      $val - the value vector to write
+    The `triton_gen.sub_group_block_write` writes a scalar or vector for each
+    work-item in the sub-group from pointer `ptr` as a block operation.
+    The data is read strided, so the first value is written to:
+    ```
+    ptr[sub_group_local_id]
+    ```
+    and the second one is:
+    ```
+    ptr[sub_group_local_id + sub_group_size]
+    ```
+    etc.
+
+    `ptr` must be aligned to the size of the element type of `res`.
+
+     Example:
+     ```mlir
+     %0 = triton_gen.sub_group_block_write %ptr, %val : !llvm.ptr<1>, vector<4xi32>
+     ```
   }];
 
+  let arguments = (ins
+      Arg<TritonGEN_SubGroupBlockMemoryAccessPointerType, "", [MemRead]>:$ptr,
+      TritonGEN_SubGroupBlockMemoryAccessType:$val);
+
+  let results = (outs);
+
   let assemblyFormat = [{
-    operands ` ` attr-dict `:` `(` type(operands) `)`
+    $ptr `,` $val attr-dict `:` qualified(type($ptr)) `,` type($val)
   }];
-
-  let hasVerifier = 1;
 }
+
 #endif // TRITONGEN_OPS

third_party/intel/lib/Dialect/TritonGEN/IR/TritonGENOps.cpp

@@ -48,18 +48,6 @@ template <typename Op> static LogicalResult verifyMatrixInput(Op op) {
   return success();
 }
 
-static LogicalResult verifySIMDBlockTy(Operation *op, VectorType vecTy) {
-  unsigned numElems = vecTy.getNumElements();
-  IntegerType elemTy = cast<IntegerType>(vecTy.getElementType());
-
-  // FIXME: Allow 16xi16 when SPIRV-LLVM translator supports it.
-  if (numElems != 1 && numElems != 2 && numElems != 4 && numElems != 8 &&
-      (elemTy.getWidth() != 8 || numElems != 16))
-    return op->emitOpError("unsupported vector type");
-
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // gen.sub_group_reduce
 //===----------------------------------------------------------------------===//
@@ -438,19 +426,3 @@ LogicalResult TritonGEN::Matrix2DBlockPrefetchOp::verify() {
 
   return success();
 }
-
-//===----------------------------------------------------------------------===//
-// gen.simdblockread
-//===----------------------------------------------------------------------===//
-
-LogicalResult TritonGEN::SIMDBlockReadOp::verify() {
-  return verifySIMDBlockTy(*this, getRes().getType());
-}
-
-//===----------------------------------------------------------------------===//
-// gen.simdblockwrite
-//===----------------------------------------------------------------------===//
-
-LogicalResult TritonGEN::SIMDBlockWriteOp::verify() {
-  return verifySIMDBlockTy(*this, getVal().getType());
-}