Moves pytest, adds pytest to CI, verifies for static input args to vi…

…ew_slice

.github/workflows/integration-tests.yml

@@ -378,9 +378,10 @@ jobs:
         run: |
           SHARED_LIB_DIR="${GITHUB_WORKSPACE}/python/triton/_C"
           if [ ! -d "${SHARED_LIB_DIR}" ]; then
-            echo "Coult not find '${SHARED_LIB_DIR}'" ; exit -1
+            echo "Could not find '${SHARED_LIB_DIR}'" ; exit -1
           fi
           pytest --capture=tee-sys -rfs python/tutorials/06-fused-attention.py
+          pytest --capture=tee-sys -rfs third_party/amd/python/test/test_view_slice.py
           cd python/test/unit
           pytest --capture=tee-sys -rfs -n 16 language runtime \
                  --ignore=language/test_line_info.py

.github/workflows/integration-tests.yml.in

@@ -386,9 +386,10 @@ jobs:
         run: |
           SHARED_LIB_DIR="${GITHUB_WORKSPACE}/python/triton/_C"
           if [ ! -d "${SHARED_LIB_DIR}" ]; then
-            echo "Coult not find '${SHARED_LIB_DIR}'" ; exit -1
+            echo "Could not find '${SHARED_LIB_DIR}'" ; exit -1
           fi
           pytest --capture=tee-sys -rfs python/tutorials/06-fused-attention.py
+          pytest --capture=tee-sys -rfs third_party/amd/python/test/test_view_slice.py
           cd python/test/unit
           pytest --capture=tee-sys -rfs -n 16 language runtime \
                  --ignore=language/test_line_info.py

test/Conversion/amd/invalid_viewslice_to_llvm.mlir

@@ -68,3 +68,33 @@ tt.func @invalid_stride(%arg0: tensor<256x128xi32, #blocked1> {tt.divisibility =
   %1 = amdgpu.view_slice %arg0[0,0] [256, 16] [1,2] : tensor<256x128xi32, #blocked1> to tensor<256x16xi32, #blocked1>
   tt.return
 }
+
+// -----
+
+// Invalid dynamic offset
+#blocked1 = #triton_gpu.blocked<{sizePerThread = [8, 1], threadsPerWarp = [4, 16], warpsPerCTA = [8, 1], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [0, 1]}>
+tt.func @invalid_stride(%arg0: tensor<256x128xi32, #blocked1> {tt.divisibility = 16 : i32}, %arg1: i32) {
+  // expected-error @+1 {{currently only static offsets are supported}}
+  %2 = amdgpu.view_slice %arg0[0,%arg1] [256, 16] [1,1] : tensor<256x128xi32, #blocked1> to tensor<256x16xi32, #blocked1>
+  tt.return
+}
+
+// -----
+
+// Invalid dynamic size
+#blocked1 = #triton_gpu.blocked<{sizePerThread = [8, 1], threadsPerWarp = [4, 16], warpsPerCTA = [8, 1], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [0, 1]}>
+tt.func @invalid_stride(%arg0: tensor<256x128xi32, #blocked1> {tt.divisibility = 16 : i32}, %arg1: i32) {
+  // expected-error @+1 {{currently only static sizes are supported}}
+  %2 = amdgpu.view_slice %arg0[0,0] [256, %arg1] [1,1] : tensor<256x128xi32, #blocked1> to tensor<256x16xi32, #blocked1>
+  tt.return
+}
+
+// -----
+
+// Invalid dynamic stride
+#blocked1 = #triton_gpu.blocked<{sizePerThread = [8, 1], threadsPerWarp = [4, 16], warpsPerCTA = [8, 1], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [0, 1]}>
+tt.func @invalid_stride(%arg0: tensor<256x128xi32, #blocked1> {tt.divisibility = 16 : i32}, %arg1: i32) {
+  // expected-error @+1 {{currently only static strides are supported}}
+  %2 = amdgpu.view_slice %arg0[0,0] [256, 16] [1,%arg1] : tensor<256x128xi32, #blocked1> to tensor<256x16xi32, #blocked1>
+  tt.return
+}

third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td

@@ -33,15 +33,15 @@ include "mlir/Interfaces/ViewLikeInterface.td" // OffsetSizeAndStrideOpInterface
 include "TritonAMDGPUDialect.td"
 include "TritonAMDGPUAttrDefs.td"
 
-class TritonAMDGPU_Op<string mnemonic, list<Trait> traits = []>
+class TT_AMDGPU_Op<string mnemonic, list<Trait> traits = []>
     : Op<TritonAMDGPU_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // ViewSliceOp
 //===----------------------------------------------------------------------===//
 
-def TritonAMDGPU_ViewSliceOp
-    : TritonAMDGPU_Op<"view_slice", [AttrSizedOperandSegments,
+def ViewSliceOp
+    : TT_AMDGPU_Op<"view_slice", [AttrSizedOperandSegments,
                                      OffsetSizeAndStrideOpInterface, Pure]> {
   let summary = "view slice operation";
   let description = [{

third_party/amd/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp

@@ -76,6 +76,20 @@ LogicalResult ViewSliceOp::verify() {
   shapePerCTA[0] = std::min(static_cast<unsigned>(srcShape[0]), shapePerCTA[0]);
   shapePerCTA[1] = std::min(static_cast<unsigned>(srcShape[1]), shapePerCTA[1]);
 
+  auto checkForConstInts = [](Value val) {
+    return isa<arith::ConstantIntOp>(val.getDefiningOp());
+  };
+
+  if (!llvm::all_of(getOffsets(), checkForConstInts)) {
+    return emitError("currently only static offsets are supported");
+  }
+  if (!llvm::all_of(getSizes(), checkForConstInts)) {
+    return emitError("currently only static sizes are supported");
+  }
+  if (!llvm::all_of(getStrides(), checkForConstInts)) {
+    return emitError("currently only static strides are supported");
+  }
+
   auto offsets = getStaticOffsets();
   auto sizes = getStaticSizes();
 

third_party/amd/python/test/test_core.py → third_party/amd/python/test/test_view_slice.py

@@ -1,4 +1,3 @@
-# flake8: noqa: F821,F841
 import tempfile
 
 import numpy as np