address review comments

Author: alexbaden

test/TritonIntelGPU/tritonintelgpu-invalid.mlir

@@ -163,22 +163,17 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttg.thr
 // -----
 
 #dpas = #ttig.dpas<{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 2, threadsPerWarp = 16, warpsPerCTA = [2, 2], repCluster = [1, 1]}>
-#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
-#blocked = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [1, 4], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
 #dot_operand_a = #ttg.dot_op<{opIdx=0, parent=#dpas, kWidth=1}>
 #dot_operand_b = #ttg.dot_op<{opIdx=1, parent=#dpas, kWidth=2}>
 #smem = #ttg.shared_memory
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: matmul_tf32dot
-  tt.func @matmul_tf32dot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},
-  %a:!ttg.memdesc<32x16xf32, #shared, #smem>, %b:!ttg.memdesc<16x32xf32, #shared, #smem>) {
+  tt.func @matmul_tf32dot(%ptr:!tt.ptr<f32>,
+  %a_mat:tensor<32x16xf32, #dot_operand_a>, %b_mat:tensor<16x32xf32, #dot_operand_b>) {
     %cst = arith.constant dense<0.000000e+00> : tensor<32x32xf32, #dpas>
-    %a_mat = ttg.local_load %a : !ttg.memdesc<32x16xf32, #shared, #smem> -> tensor<32x16xf32, #dot_operand_a>
-    %b_mat = ttg.local_load %b : !ttg.memdesc<16x32xf32, #shared, #smem> -> tensor<16x32xf32, #dot_operand_b>
 
     // expected-error @+1 {{Layout has opsPerChannel = 2 but tensor element type is 'f32'. Expected 16 bit type.}}
     %28 = tt.dot %a_mat, %b_mat, %cst, inputPrecision = tf32 : tensor<32x16xf32, #dot_operand_a> * tensor<16x32xf32, #dot_operand_b> -> tensor<32x32xf32, #dpas>
-    %38 = ttg.convert_layout %28 : tensor<32x32xf32, #dpas> -> tensor<32x32xf32, #blocked>
 
     tt.return
   }
@@ -187,22 +182,16 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 // -----
 
 #dpas = #ttig.dpas<{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 1, threadsPerWarp = 16, warpsPerCTA = [2, 2], repCluster = [1, 1]}>
-#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
-#blocked = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [1, 4], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
 #dot_operand_a = #ttg.dot_op<{opIdx=0, parent=#dpas, kWidth=1}>
 // expected-error @below {{ttg.dot_op kWidth parameter must match the parent's opsPerChannel}}
 #dot_operand_b = #ttg.dot_op<{opIdx=1, parent=#dpas, kWidth=2}>
 #smem = #ttg.shared_memory
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: matmul_tf32dot
-  tt.func @matmul_tf32dot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},
-  %a:!ttg.memdesc<32x16xf32, #shared, #smem>, %b:!ttg.memdesc<16x32xf32, #shared, #smem>) {
+  tt.func @matmul_tf32dot(%ptr:!tt.ptr<f32>,
+  %a_mat:tensor<32x16xf32, #dot_operand_a>, %b_mat:tensor<16x32xf32, #dot_operand_b>) {
     %cst = arith.constant dense<0.000000e+00> : tensor<32x32xf32, #dpas>
-    %a_mat = ttg.local_load %a : !ttg.memdesc<32x16xf32, #shared, #smem> -> tensor<32x16xf32, #dot_operand_a>
-    %b_mat = ttg.local_load %b : !ttg.memdesc<16x32xf32, #shared, #smem> -> tensor<16x32xf32, #dot_operand_b>
-
     %28 = tt.dot %a_mat, %b_mat, %cst, inputPrecision = tf32 : tensor<32x16xf32, #dot_operand_a> * tensor<16x32xf32, #dot_operand_b> -> tensor<32x32xf32, #dpas>
-    %38 = ttg.convert_layout %28 : tensor<32x32xf32, #dpas> -> tensor<32x32xf32, #blocked>
 
     tt.return
   }

third_party/intel/lib/Dialect/TritonIntelGPU/IR/Dialect.cpp

@@ -1222,40 +1222,30 @@ struct TritonIntelGPUVerifyTensorLayoutInterface
         if (auto ptrTy = dyn_cast<PointerType>(elemTy)) {
           elemTy = ptrTy.getPointeeType();
         }
-        const auto elemTyBitWidth = elemTy.getIntOrFloatBitWidth();
+        const unsigned elemTyBitWidth = elemTy.getIntOrFloatBitWidth();
 
         // We know opsPerChannel is either 1, 4, or 8 because of the DPAS
         // verifier when the DPAS attribute is created. Here we verify that
         // opsPerChannel matches the tensor type.
-        if (dpasEncoding.getOpsPerChannel() == 4 && elemTyBitWidth != 8) {
+        if (dpasEncoding.getOpsPerChannel() * elemTyBitWidth != 32) {
           return makeErr() << layout << ".\nLayout has opsPerChannel = "
                            << dpasEncoding.getOpsPerChannel()
                            << " but tensor element type is " << elemTy
-                           << ". Expected 8 bit type.";
-        } else if (dpasEncoding.getOpsPerChannel() == 2 &&
-                   elemTyBitWidth != 16) {
-          return makeErr() << layout << ".\nLayout has opsPerChannel = "
-                           << dpasEncoding.getOpsPerChannel()
-                           << " but tensor element type is " << elemTy
-                           << ". Expected 16 bit type.";
-        } else if (dpasEncoding.getOpsPerChannel() == 1 &&
-                   elemTyBitWidth != 32) {
-          return makeErr() << layout << ".\nLayout has opsPerChannel = "
-                           << dpasEncoding.getOpsPerChannel()
-                           << " but tensor element type is " << elemTy
-                           << ". Expected 32 bit type.";
+                           << ". Expected "
+                           << 32 / dpasEncoding.getOpsPerChannel()
+                           << " bit type.";
         }
         return success();
       };
 
-      if (isa<DotOp>(op)) {
-        auto dotOp = cast<DotOp>(op);
+      if (auto dotOp = dyn_cast<DotOp>(op)) {
         auto aElemTy = dotOp.getA().getType().getElementType();
-        auto result = validateDotDpasLayout(aElemTy);
-        if (result.failed())
-          return result;
-
         auto bElemTy = dotOp.getB().getType().getElementType();
+
+        auto aResult = validateDotDpasLayout(aElemTy);
+        if (aResult.failed())
+          return aResult;
+
         return validateDotDpasLayout(bElemTy);
       }
     }