[NVPTX, CUDA] Add .and.popc variant of the b1 MMA instruction.

[tra]

Summary

Extends the changes in D104847 and adds another MMA instruction variant and corresponding intrinsics & builtins.

That should allow clang to compile mma.h from CUDA-11.3.

Didn't test it much yet. There may still be some sharp corners.

[steffenlarsen]

Good stuff! Thanks for adding this and adjusting the test generator. I have requested some minor changes, though nothing critical. Are the test failures related to these changes?

clang/test/CodeGen/builtins-nvptx-mma.py

35–39

84

The following changes would remove the need for the m16n16k8 cases in is_ldst_variant_supported.

261

[tra]

Addressed review comments.

[tra]

clang/test/CodeGen/builtins-nvptx-mma.py

35–39

Default initializers that use objects in python are one of the common gotchas.
https://docs.python-guide.org/writing/gotchas/#mutable-default-arguments

It probably does not make much of a difference in this case as we do not modify it, but I'd prefer to avoid it nevertheless.

84

This was done deliberately. I did have that in an earlier variant of the patch, but eventually settled on the current version.

My thinking is that make_ldst_ops(m16n16k8) would create whatever is necessary for m16n16k8, and we'd keep the decision of what to produce in one place in is_ldst_variant_supported. Splitting one make_ldst_ops into two here makes this decision implicit which would need additional explanations. Code that needs less explanations wins. :-)

[steffenlarsen]

Thank you for addressing my concerns. I am happy with the changes. Great work!

clang/test/CodeGen/builtins-nvptx-mma.py

35–39

Interesting and a little horrifying. I agree it wouldn't have been a problem in this particular case, but I understand your concern and I am fine with keeping it as-is. 👍

84

My concern is that it is more confusing this way because they are the only load/store ops generated where the type is then later filtered. It makes sense for MMA because it needs to filter out select combinations, but here the comment above says one thing and make_ldst_ops does something else.

I don't think it makes a huge difference either way, so if you prefer the current version I am okay with keeping it. That said, it might be good to have a comment in the m16n16k8 case of is_ldst_variant_supported making the connection to this.

[tra]

Updated LD/ST generation.

[tra]

clang/test/CodeGen/builtins-nvptx-mma.py

84

Fair enough. I've adopted your suggestion. Builtins don't have as much magic as intrinsics, so your approach is indeed simpler.

[steffenlarsen]

LGTM! Are the test failures related to this, you reckon?

[tra]

rebase

[tra]

In D105384#2879308, @steffenlarsen wrote:

LGTM! Are the test failures related to this, you reckon?

AFAICT, no, the test failures don't seem to be related. It appears that the test runs rarely succeed in general. They are mostly red for nearly all build attempts (https://reviews.llvm.org/harbormaster/build/?plan=PHID-HMCP-p2oc4ocen3l2yzymvg2l), even though LLVM buildbots are green (https://lab.llvm.org/buildbot/#/).

clang/include/clang/Basic/BuiltinsNVPTX.def

 TARGET_BUILTIN(__hmma_m8n32k16_mma_f32f16, "vf*iC*iC*iC*IiIi", "", AND(SM_70,PTX61))
 TARGET_BUILTIN(__hmma_m8n32k16_mma_f32f32, "vf*iC*iC*fC*IiIi", "", AND(SM_70,PTX61))
 TARGET_BUILTIN(__hmma_m8n32k16_mma_f16f32, "vi*iC*iC*fC*IiIi", "", AND(SM_70,PTX61))
 
 // Builtins to support integer and sub-integer WMMA instructions on sm_72/sm_75
 TARGET_BUILTIN(__bmma_m8n8k128_ld_a_b1, "vi*iC*UiIi", "", AND(SM_75,PTX63))
 TARGET_BUILTIN(__bmma_m8n8k128_ld_b_b1, "vi*iC*UiIi", "", AND(SM_75,PTX63))
 TARGET_BUILTIN(__bmma_m8n8k128_ld_c, "vi*iC*UiIi", "", AND(SM_75,PTX63))
+TARGET_BUILTIN(__bmma_m8n8k128_mma_and_popc_b1, "vi*iC*iC*iC*Ii", "", AND(SM_75,PTX71))
 TARGET_BUILTIN(__bmma_m8n8k128_mma_xor_popc_b1, "vi*iC*iC*iC*Ii", "", AND(SM_75,PTX63))
 TARGET_BUILTIN(__bmma_m8n8k128_st_c_i32, "vi*iC*UiIi", "", AND(SM_75,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_ld_a_s8, "vi*iC*UiIi", "", AND(SM_72,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_ld_a_u8, "vi*iC*UiIi", "", AND(SM_72,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_ld_b_s8, "vi*iC*UiIi", "", AND(SM_72,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_ld_b_u8, "vi*iC*UiIi", "", AND(SM_72,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_ld_c, "vi*iC*UiIi", "", AND(SM_72,PTX63))
 TARGET_BUILTIN(__imma_m16n16k16_mma_s8, "vi*iC*iC*iC*IiIi", "", AND(SM_72,PTX63))

clang/lib/CodeGen/CGBuiltin.cpp

#define MMA_VARIANTS_I4(geom, type) \
       Intrinsic::nvvm_wmma_##geom##_mma_row_col_##type,             \
       0, \
       0, \
       0, \
       Intrinsic::nvvm_wmma_##geom##_mma_row_col_##type##_satfinite, \
       0, \
       0
 // b1 MMA does not support .satfinite.
-#define MMA_VARIANTS_B1(geom, type) \
+#define MMA_VARIANTS_B1_XOR(geom, type) \
       0, \
-      Intrinsic::nvvm_wmma_##geom##_mma_row_col_##type,             \
+      Intrinsic::nvvm_wmma_##geom##_mma_xor_popc_row_col_##type,             \
+      0, \
+      0, \
+      0, \
+      0, \
+      0, \
+      0
+#define MMA_VARIANTS_B1_AND(geom, type) \
+      0, \
+      Intrinsic::nvvm_wmma_##geom##_mma_and_popc_row_col_##type,             \
       0, \
       0, \
       0, \
       0, \
       0, \
       0
   // clang-format on
   switch (BuiltinID) {
case NVPTX::BI__imma_m8n32k16_mma_u8:
     return {1, 4, 8, 8, {{MMA_SATF_VARIANTS(m8n32k16, u8)}}};
 
   // Sub-integer MMA
   case NVPTX::BI__imma_m8n8k32_mma_s4:
     return {1, 1, 2, 2, {{MMA_VARIANTS_I4(m8n8k32, s4)}}};
   case NVPTX::BI__imma_m8n8k32_mma_u4:
     return {1, 1, 2, 2, {{MMA_VARIANTS_I4(m8n8k32, u4)}}};
   case NVPTX::BI__bmma_m8n8k128_mma_xor_popc_b1:
-    return {1, 1, 2, 2, {{MMA_VARIANTS_B1(m8n8k128, b1)}}};
+    return {1, 1, 2, 2, {{MMA_VARIANTS_B1_XOR(m8n8k128, b1)}}};
+  case NVPTX::BI__bmma_m8n8k128_mma_and_popc_b1:
+    return {1, 1, 2, 2, {{MMA_VARIANTS_B1_AND(m8n8k128, b1)}}};
 
   // Double MMA
   case NVPTX::BI__dmma_m8n8k4_mma_f64:
     return {1, 1, 2, 2, {{MMA_VARIANTS(m8n8k4, f64)}}};
 
   // Alternate FP MMA
   case NVPTX::BI__mma_bf16_m16n16k16_mma_f32:
     return {4, 4, 8, 8, {{MMA_VARIANTS(m16n16k16, bf16)}}};
   case NVPTX::BI__mma_bf16_m8n32k16_mma_f32:
     return {2, 8, 8, 8, {{MMA_VARIANTS(m8n32k16, bf16)}}};
   case NVPTX::BI__mma_bf16_m32n8k16_mma_f32:
     return {8, 2, 8, 8, {{MMA_VARIANTS(m32n8k16, bf16)}}};
   case NVPTX::BI__mma_tf32_m16n16k8_mma_f32:
     return {4, 4, 8, 8, {{MMA_VARIANTS(m16n16k8, tf32)}}};
   default:
     llvm_unreachable("Unexpected builtin ID.");
   }
 #undef MMA_VARIANTS
 #undef MMA_SATF_VARIANTS
 #undef MMA_VARIANTS_I4
-#undef MMA_VARIANTS_B1
+#undef MMA_VARIANTS_B1_AND
+#undef MMA_VARIANTS_B1_XOR
 }
 
 } // namespace
 
 Value *
 CodeGenFunction::EmitNVPTXBuiltinExpr(unsigned BuiltinID, const CallExpr *E) {
   auto MakeLdg = [&](unsigned IntrinsicID) {
     Value *Ptr = EmitScalarExpr(E->getArg(0));
CodeGenFunction::EmitNVPTXBuiltinExpr(unsigned BuiltinID, const CallExpr *E) {
   case NVPTX::BI__imma_m16n16k16_mma_u8:
   case NVPTX::BI__imma_m32n8k16_mma_s8:
   case NVPTX::BI__imma_m32n8k16_mma_u8:
   case NVPTX::BI__imma_m8n32k16_mma_s8:
   case NVPTX::BI__imma_m8n32k16_mma_u8:
   case NVPTX::BI__imma_m8n8k32_mma_s4:
   case NVPTX::BI__imma_m8n8k32_mma_u4:
   case NVPTX::BI__bmma_m8n8k128_mma_xor_popc_b1:
+  case NVPTX::BI__bmma_m8n8k128_mma_and_popc_b1:
   case NVPTX::BI__dmma_m8n8k4_mma_f64:
   case NVPTX::BI__mma_bf16_m16n16k16_mma_f32:
   case NVPTX::BI__mma_bf16_m8n32k16_mma_f32:
   case NVPTX::BI__mma_bf16_m32n8k16_mma_f32:
   case NVPTX::BI__mma_tf32_m16n16k8_mma_f32: {
     Address Dst = EmitPointerWithAlignment(E->getArg(0));
     Address SrcA = EmitPointerWithAlignment(E->getArg(1));
     Address SrcB = EmitPointerWithAlignment(E->getArg(2));
     Address SrcC = EmitPointerWithAlignment(E->getArg(3));
     Optional<llvm::APSInt> LayoutArg =
         E->getArg(4)->getIntegerConstantExpr(getContext());
     if (!LayoutArg)
       return nullptr;
     int Layout = LayoutArg->getSExtValue();
     if (Layout < 0 || Layout > 3)
       return nullptr;
     llvm::APSInt SatfArg;
-    if (BuiltinID == NVPTX::BI__bmma_m8n8k128_mma_xor_popc_b1)
+    if (BuiltinID == NVPTX::BI__bmma_m8n8k128_mma_xor_popc_b1 ||
+        BuiltinID == NVPTX::BI__bmma_m8n8k128_mma_and_popc_b1)
       SatfArg = 0;  // .b1 does not have satf argument.
     else if (Optional<llvm::APSInt> OptSatfArg =
                  E->getArg(5)->getIntegerConstantExpr(getContext()))
       SatfArg = *OptSatfArg;
     else
       return nullptr;
     bool Satf = SatfArg.getSExtValue();
     NVPTXMmaInfo MI = getNVPTXMmaInfo(BuiltinID);

clang/test/CodeGen/builtins-nvptx-mma.cu

 
 //
 // *** DO NOT EDIT ***
 //
 //  This test has been automatically generated by
-//  builtins-nvtx-mma.py --ptx=70 --gpu-arch=80
+//  builtins-nvtx-mma.py --ptx=71 --gpu-arch=80
 //
-// Make sure we can handle all builtins available on sm_80 with PTX70
+// Make sure we can handle all builtins available on sm_80 with PTX71
 // RUN: %clang_cc1 -triple nvptx64-unknown-unknown -target-cpu sm_80 \
-// RUN:            -fcuda-is-device -target-feature +ptx70 \
-// RUN:            -DPTX=70 -DSM=80 \
+// RUN:            -fcuda-is-device -target-feature +ptx71 \
+// RUN:            -DPTX=71 -DSM=80 \
 // RUN:            -S -emit-llvm -o - -x cuda %s \
-// RUN:   | FileCheck -check-prefixes=CHECK_PTX70_SM80,CHECK_PTX60_SM70,CHECK_PTX63_SM72,CHECK_PTX61_SM70,CHECK_PTX63_SM75 %s
+// RUN:   | FileCheck -check-prefixes=CHECK_PTX70_SM80,CHECK_PTX60_SM70,CHECK_PTX63_SM72,CHECK_PTX61_SM70,CHECK_PTX63_SM75,CHECK_PTX71_SM75 %s
 // Verify that all builtins have correct constraints.
 // RUN: %clang_cc1 -triple nvptx-unknown-unknown \
 // RUN:   -target-cpu sm_60 -target-feature +ptx42 \
-// RUN:   -DPTX=70 -DSM=80 -fcuda-is-device -S -o /dev/null -x cuda \
+// RUN:   -DPTX=71 -DSM=80 -fcuda-is-device -S -o /dev/null -x cuda \
 // RUN:   -verify %s
 
+
 #if !defined(CUDA_VERSION)
 #define __device__ __attribute__((device))
 #define __global__ __attribute__((global))
 #define __shared__ __attribute__((shared))
 #define __constant__ __attribute__((constant))
 
 typedef unsigned long long uint64_t;
 #endif
 
 // CHECK-LABEL: test_wmma_buitins
 __device__ void test_wmma_buitins(int *src, int *dst,
                                   float *fsrc, float *fdst,
                                   double *dsrc, double *ddst, int ldm) {
 
+
 #if (PTX >= 60) && (SM >= 70)
 
   // CHECK_PTX60_SM70: call {{.*}} @llvm.nvvm.wmma.m16n16k16.load.a.col.stride.f16
   // expected-error-re@+1 {{'__hmma_m16n16k16_ld_a' needs target feature (sm_70{{.*}},(ptx60{{.*}}}}
   __hmma_m16n16k16_ld_a(dst, src, ldm, 1);
   // CHECK_PTX60_SM70: call {{.*}} @llvm.nvvm.wmma.m16n16k16.load.a.row.stride.f16
   // expected-error-re@+1 {{'__hmma_m16n16k16_ld_a' needs target feature (sm_70{{.*}},(ptx60{{.*}}}}
   __hmma_m16n16k16_ld_a(dst, src, ldm, 0);
#if (PTX >= 63) && (SM >= 75)
   // expected-error-re@+1 {{'__imma_m8n8k32_ld_c' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_ld_c(dst, src, ldm, 0);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.store.d.col.stride.s32
   // expected-error-re@+1 {{'__imma_m8n8k32_st_c_i32' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_st_c_i32(dst, src, ldm, 1);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.store.d.row.stride.s32
   // expected-error-re@+1 {{'__imma_m8n8k32_st_c_i32' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_st_c_i32(dst, src, ldm, 0);
-  // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k128.mma.row.col.b1
+  // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k128.mma.xor.popc.row.col.b1
   // expected-error-re@+1 {{'__bmma_m8n8k128_mma_xor_popc_b1' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __bmma_m8n8k128_mma_xor_popc_b1(dst, src, src, src, 1);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.mma.row.col.s4
   // expected-error-re@+1 {{'__imma_m8n8k32_mma_s4' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_mma_s4(dst, src, src, src, 1, 0);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.mma.row.col.s4.satfinite
   // expected-error-re@+1 {{'__imma_m8n8k32_mma_s4' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_mma_s4(dst, src, src, src, 1, 1);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.mma.row.col.u4
   // expected-error-re@+1 {{'__imma_m8n8k32_mma_u4' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_mma_u4(dst, src, src, src, 1, 0);
   // CHECK_PTX63_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k32.mma.row.col.u4.satfinite
   // expected-error-re@+1 {{'__imma_m8n8k32_mma_u4' needs target feature (sm_75{{.*}},(ptx63{{.*}}}}
   __imma_m8n8k32_mma_u4(dst, src, src, src, 1, 1);
-#endif // (PTX >= 63) && (SM >= 75)
+#endif // (PTX >= 63) && (SM >= 75) 
 
 #if (PTX >= 70) && (SM >= 80)
 
   // CHECK_PTX70_SM80: call {{.*}} @llvm.nvvm.wmma.m16n16k16.load.a.col.stride.bf16
   // expected-error-re@+1 {{'__mma_bf16_m16n16k16_ld_a' needs target feature (sm_80{{.*}},(ptx70{{.*}}}}
   __mma_bf16_m16n16k16_ld_a(dst, src, ldm, 1);
   // CHECK_PTX70_SM80: call {{.*}} @llvm.nvvm.wmma.m16n16k16.load.a.row.stride.bf16
   // expected-error-re@+1 {{'__mma_bf16_m16n16k16_ld_a' needs target feature (sm_80{{.*}},(ptx70{{.*}}}}
#if (PTX >= 70) && (SM >= 80)
   // expected-error-re@+1 {{'__dmma_m8n8k4_mma_f64' needs target feature (sm_80{{.*}},(ptx70{{.*}}}}
   __dmma_m8n8k4_mma_f64(ddst, dsrc, dsrc, dsrc, 2, 0);
   // CHECK_PTX70_SM80: call {{.*}} @llvm.nvvm.wmma.m8n8k4.mma.row.col.f64
   // expected-error-re@+1 {{'__dmma_m8n8k4_mma_f64' needs target feature (sm_80{{.*}},(ptx70{{.*}}}}
   __dmma_m8n8k4_mma_f64(ddst, dsrc, dsrc, dsrc, 1, 0);
   // CHECK_PTX70_SM80: call {{.*}} @llvm.nvvm.wmma.m8n8k4.mma.row.row.f64
   // expected-error-re@+1 {{'__dmma_m8n8k4_mma_f64' needs target feature (sm_80{{.*}},(ptx70{{.*}}}}
   __dmma_m8n8k4_mma_f64(ddst, dsrc, dsrc, dsrc, 0, 0);
-#endif // (PTX >= 70) && (SM >= 80)
+#endif // (PTX >= 70) && (SM >= 80) 
+
+#if (PTX >= 71) && (SM >= 75)
+
+  // CHECK_PTX71_SM75: call {{.*}} @llvm.nvvm.wmma.m8n8k128.mma.and.popc.row.col.b1
+  // expected-error-re@+1 {{'__bmma_m8n8k128_mma_and_popc_b1' needs target feature (sm_75{{.*}},(ptx71{{.*}}}}
+  __bmma_m8n8k128_mma_and_popc_b1(dst, src, src, src, 1);
+#endif // (PTX >= 71) && (SM >= 75) 
 }

clang/test/CodeGen/builtins-nvptx-mma.py

def __init__(self, geom, frag, ptx_elt_type):
     self.geom = geom
     self.frag = frag
     self.ptx_type = ptx_elt_type;
 
   def __repr__(self):
     return "%s:%s:%s" % (self.geom, self.frag, self.ptx_type)
 
 class MMAOp:
-  def __init__(self, a, b, c, d):
+  def __init__(self, a, b, c, d, b1op=""):
     self.a = a
     self.b = b
     self.c = c
     self.d = d
+    self.b1op = b1op
 
   def __repr__(self):
     return ("{A:%s, B:%s, C:%s, D:%s}" % (self.a, self.b, self.c, self.d ))
 
-def make_mma_ops(geoms, types_a, types_b, types_c, types_d):
+def make_mma_ops(geoms, types_a, types_b, types_c, types_d, b1ops=None):
   ops = []
+  if b1ops is None:
+    b1ops = [""]
   for geom, type_a, type_c in product( geoms,  types_a, types_c):
     for type_b, type_d in product(types_b if types_b else [type_a],
                                   types_d if types_d else [type_c]):
-      ops.append(MMAOp(MMAFrag(geom, "a", type_a),
-                       MMAFrag(geom, "b", type_b),
-                       MMAFrag(geom, "c", type_c),
-                       MMAFrag(geom, "d", type_d)))
+      ops += [
+          MMAOp(MMAFrag(geom, "a", type_a),
+                MMAFrag(geom, "d", type_d), b1op)
+          for b1op in b1ops]
   return ops
 
 def make_ldst_ops(geoms, frags, types):
   return [MMAFrag(geom, frag, ptx_type) for (geom, frag, ptx_type)
           in product(geoms, frags, types)]
 
 def get_mma_ops():
   return (make_mma_ops(["m16n16k8"],
                        ["tf32"], [], ["f32"], []) +
           make_mma_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
                        ["bf16"], [], ["f32"], []) +
           make_mma_ops(["m8n8k4"],
                        ["f64"], [], ["f64"], []) +
           make_mma_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
                        ["f16"], [], ["f16", "f32"], ["f16", "f32"]) +
           make_mma_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
                        ["s8", "u8"], [], ["s32"], []) +
           make_mma_ops(["m8n8k32"],
                        ["s4", "u4"], [], ["s32"], []) +
           make_mma_ops(["m8n8k128"],
-                       ["b1"], [], ["s32"], []))
+                       ["b1"], [], ["s32"], [],
+                       [".xor.popc", ".and.popc"]))
 
 def get_ldst_ops():
+  # NOTE: fragemts are from the point of view of PTX.
+  # fragment `d` is only for store ops, others for both loads and stores.
   return (make_ldst_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
                         ["a", "b"], ["f16", "u8", "s8", "bf16"]) +
           make_ldst_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
                         ["c", "d"], ["f16", "f32", "s32"]) +
           make_ldst_ops(["m8n8k32"], ["a", "b"], ["s4","u4"]) +
           make_ldst_ops(["m8n8k128"], ["a", "b"], ["b1"]) +
           make_ldst_ops(["m8n8k32", "m8n8k128"],  ["c", "d"], ["s32"]) +
           make_ldst_ops(["m8n8k4"], ["a", "b", "c", "d"], ["f64"]) +
-          make_ldst_ops(["m16n16k8"], ["a", "b"], ["tf32"]) +
-          make_ldst_ops(["m16n16k8"], ["c", "d"], ["f32"]))
+          # TF32 m16n16k8 is odd.
+          # For fragment 'C' it uses __mma_*tf32*_m16n16k8_ld_c
+          # but 'D' calls __mma_m16n16k8_st_c_*f32*.
+          make_ldst_ops(["m16n16k8"], ["a", "b", "c"], ["tf32"]) +
+          make_ldst_ops(["m16n16k8"], ["d"], ["f32"]))
 
 def is_geom_supported(geom):
   # geometries for FP and ints.
   if geom in ["m8n32k16", "m32n8k16"]:
     return ptx_version >= 61
   # geometries for sub-ints.
   if geom in ["m8n8k32", "m8n8k128"]:
     return ptx_version >= 63 and gpu_arch >= 75
def get_mma_builtin_name(op):
 
   if prefix == "__hmma":
     suffix = op.d.ptx_type + op.c.ptx_type
   elif prefix in ["__mma_bf16", "__mma_tf32"]:
     suffix = op.d.ptx_type
   else:
     suffix = op.a.ptx_type
 
-  name = "%s_%s_mma%s_%s" % (prefix, op.a.geom,
-                             "_xor_popc" if op.a.ptx_type == "b1" else "",
-                             suffix)
+  name = "{prefix}_{geom}_mma{b1op}_{suffix}".format(
+      prefix = prefix,
+      geom = op.a.geom,
+      b1op = op.b1op.replace(".","_"),
+      suffix = suffix)
   return name
 
-def get_required_sm(frag):
+def get_required_sm(frag, b1op=""):
   if frag.ptx_type in ["f64", "bf16", "tf32"]:
     return 80
   if frag.ptx_type in ["u4", "s4", "b1"]:
+    if b1op == "_and_popc":
+      return 80
     return 75
   if frag.ptx_type in ["s8", "u8"]:
     return 72
   if frag.ptx_type == "s32":
     if frag.geom in ["m8n8k32", "m8n8k128"]: # s4/u4/b1
       return 75
     else:                       # s8/u8
       return 72
   if frag.ptx_type in ["f16", "f32"]:
     if frag.geom == "m16n16k8":
       return 80
     else:
       return 70
   assert(False)
 
-def get_required_ptx(frag):
+def get_required_ptx(frag, b1op=""):
+  if frag.ptx_type == "b1" and b1op == ".and.popc":
+    return 71
   if frag.ptx_type in ["f64", "bf16", "tf32"]:
     return 70
   if frag.ptx_type in ["f16", "f32"]:
     if frag.geom == "m16n16k16":
       return 60
     if frag.geom == "m16n16k8":
       return 70
     return 61
   return 63
 
-def get_src_dst_prefix(ptx_type):
-  if ptx_type == "f32":
+def get_src_dst_prefix(frag):
+  if frag.ptx_type == "f32":
     return "f"
-  if ptx_type == "f64":
+  if frag.ptx_type == "f64":
     return "d"
+  if frag.ptx_type == "tf32" and frag.frag in ["c", "d"]:
+    return "f"
   return ""
 
 def gen_wmma_ldst_tests(results):
   load_template = """
   // CHECK${check_suffix}: call {{.*}} @${intrinsic}
   // expected-error-re@+1 {{'${builtin}' needs target feature (sm_${min_sm}{{.*}},(ptx${min_ptx}{{.*}}}}
   ${builtin}(${dst}, ${src}, ldm, ${blayout});
 """.rstrip()
   intrinsic_template = "llvm.nvvm.wmma.${geom}.${op}.${frag}.${ilayout}.stride.${itype}"
 
   for frag, layout in sorted(product(get_ldst_ops(), ["row","col"]), key=str):
 
     if not is_ldst_variant_supported(frag, layout):
       continue
 
-    src_dst_prefix = get_src_dst_prefix(frag.ptx_type)
+    src_dst_prefix = get_src_dst_prefix(frag)
+
     min_sm = get_required_sm(frag)
     min_ptx = get_required_ptx(frag)
+    # TF32 uses f32 for accumulator loads.
+    if frag.geom == "m16n16k8" and frag.frag =="c":
+      assert frag.ptx_type == "tf32"
+      itype = "f32"
+    else:
+      itype = frag.ptx_type
+
     params = {
         "check_suffix" : "_PTX%d_SM%d" % (min_ptx, min_sm),
         "builtin" : get_ldst_builtin_name(frag),
         "min_ptx" : min_ptx,
         "min_sm" : min_sm,
         "dst": src_dst_prefix + "dst",
         "src": src_dst_prefix + "src",
         "blayout" : 0 if layout == "row" else 1,
         "intrinsic" : Template(intrinsic_template).substitute({
             "frag" : frag.frag,
             "geom"   : frag.geom,
             "ilayout" : layout,
-            "itype" : frag.ptx_type,
+            "itype" : itype,
             "op" : "store" if frag.frag == "d" else "load",
         })
     }
     results[(min_ptx,min_sm)] += Template(load_template).substitute(params)
 
   return results
 
 def mma_signature(op):
def get_ilayout(a, b):
   }[a + "." + b]
 
 def gen_wmma_mma_tests(results):
   mma_template = """
   // CHECK${check_suffix}: call {{.*}} @${intrinsic}
   // expected-error-re@+1 {{'${builtin}' needs target feature (sm_${min_sm}{{.*}},(ptx${min_ptx}{{.*}}}}
   ${builtin}(${dst}, ${asrc}, ${asrc}, ${csrc}, ${ilayout}${maybe_satf});
 """.rstrip()
-  intrinsic_template = "llvm.nvvm.wmma.${geom}.mma.${alayout}.${blayout}.${intrinsic_signature}${satf}"
+  intrinsic_template = "llvm.nvvm.wmma.${geom}.mma${b1op}.${alayout}.${blayout}.${intrinsic_signature}${satf}"
 
   for op, alayout, blayout, satf in sorted(product( get_mma_ops(),
                                                     ["row","col"],
                                                     ["row","col"],
                                                     [".satfinite", ""]),
                                            key=str):
 
     if not is_mma_variant_supported(op, alayout, blayout, satf):
       continue
 
-    asrc_prefix = get_src_dst_prefix(op.a.ptx_type)
-    csrc_prefix = get_src_dst_prefix(op.c.ptx_type)
-    ddst_prefix = get_src_dst_prefix(op.d.ptx_type)
-    min_sm = get_required_sm(op.a)
-    min_ptx = get_required_ptx(op.a)
+    asrc_prefix = get_src_dst_prefix(op.a)
+    csrc_prefix = get_src_dst_prefix(op.c)
+    ddst_prefix = get_src_dst_prefix(op.d)
     if op.a.ptx_type == "b1": # .b1 MMA has no satf argument.
        isatf_arg = ""
     else:
        isatf_arg = ", 1" if satf else ", 0"
+    min_sm = get_required_sm(op.a, op.b1op)
+    min_ptx = get_required_ptx(op.a, op.b1op)
     params = {
         "check_suffix" : "_PTX%d_SM%d" % (min_ptx, min_sm),
         "builtin" : get_mma_builtin_name(op),
         "min_ptx" : min_ptx,
         "min_sm" : min_sm,
         "dst": ddst_prefix + "dst",
         "asrc": asrc_prefix + "src",
         "csrc": csrc_prefix + "src",
         "ilayout" : get_ilayout(alayout, blayout),
         "maybe_satf" : isatf_arg,
         "intrinsic" : Template(intrinsic_template).substitute({
             "geom"  : op.a.geom,
             "alayout" : alayout,
             "blayout" : blayout,
             "intrinsic_signature" : mma_signature(op),
             "satf"  : satf,
+            "b1op"  : op.b1op
         })
     }
     results[(min_ptx, min_sm)] += Template(mma_template).substitute(params)
 
   return results
 
 def gen_tests():
   results = gen_wmma_ldst_tests(defaultdict(str))

llvm/include/llvm/IR/IntrinsicsNVVM.td

list<WMMA_REGS> id_frags = !cond(
      !eq(A.ptx_elt_type, "f16") : [D, C],
      // other ops are identified by input types.
      !ne(A.ptx_elt_type, B.ptx_elt_type): [A, B],
      true: [A]
      );
    string ret = !foldl("", id_frags, a, b, !strconcat(a, ".", b.ptx_elt_type));
 }
 
-class WMMA_NAME<string ALayout, string BLayout, int Satfinite, string Rnd,
+class WMMA_NAME<string ALayout, string BLayout, int Satfinite, string Rnd, string b1op,
                 WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
   string signature = MMA_SIGNATURE<A, B, C, D>.ret;
   string llvm = "llvm.nvvm.wmma."
                 # A.geom
                 # ".mma"
+                # b1op
                 # "." # ALayout
                 # "." # BLayout
                 # !if(!ne(Rnd, ""), !strconcat(".", Rnd), "")
                 # signature
                 # !if(Satfinite, ".satfinite", "");
 
   string record = !subst(".", "_",
                   !subst("llvm.", "int_", llvm));
 }
 
-class MMA_NAME<string ALayout, string BLayout, int Satfinite,
+class MMA_NAME<string ALayout, string BLayout, int Satfinite, string b1op,
                WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
   string signature = MMA_SIGNATURE<A, B, C, D>.ret;
-  string llvm = "llvm.nvvm.mma."
-                # A.geom
+  string llvm = "llvm.nvvm.mma"
+                # b1op
+                # "." # A.geom
                 # "." # ALayout
                 # "." # BLayout
                 # !if(Satfinite, ".satfinite", "")
                 # signature;
   string record = !subst(".", "_",
                   !subst("llvm.", "int_", llvm));
 }
 
bit ret = !cond(
     !and(!or(!eq(t, "s4"),
              !eq(t, "u4"),
              !eq(t, "b1")),
          !ne(layout, "row:col")) : false,
     true: true
   );
 }
 
+class NVVM_MMA_B1OPS<list<WMMA_REGS> frags> {
+  list<string> ret = !cond(
+    !eq(frags[0].ptx_elt_type, "b1") : [".xor.popc", ".and.popc"],
+    true: [""]
+  );
+}
+
 // Returns true if this combination of layout/satf for MMA ops is supported;
 // false otherwise.
 // E.g.
 // if NVVM_MMA_SUPPORTED<...>.ret then
 //   def : FOO<>; // The record will only be defined for supported ops.
 //
 class NVVM_MMA_SUPPORTED<list<WMMA_REGS> frags, string layout_a, string layout_b, int satf> {
   // MMA ops check both layouts.
foreach stride = [0, 1] in {
     foreach frag = NVVM_MMA_OPS.all_st_ops in
       if NVVM_WMMA_LDST_SUPPORTED<frag, layout>.ret then
         def WMMA_NAME_LDST<"store", frag, layout, stride>.record
              : NVVM_WMMA_ST<frag, layout, stride>;
   }
 }
 
 // WMMA.MMA
-class NVVM_WMMA_MMA<string ALayout, string BLayout, int Satfinite, string rnd,
+class NVVM_WMMA_MMA<string ALayout, string BLayout, int Satfinite, string rnd, string b1op,
                     WMMA_REGS A, WMMA_REGS B,
                     WMMA_REGS C, WMMA_REGS D>
   : Intrinsic<D.regs,
               !listconcat(A.regs, B.regs, C.regs),
               [IntrNoMem],
-              WMMA_NAME<ALayout, BLayout, Satfinite, rnd, A, B, C, D>.llvm>;
+              WMMA_NAME<ALayout, BLayout, Satfinite, rnd, b1op, A, B, C, D>.llvm>;
 
 foreach layout_a = ["row", "col"] in {
   foreach layout_b = ["row", "col"] in {
     foreach satf = [0, 1] in {
       foreach rnd = ["", "rn", "rz", "rm", "rp"] in {
         foreach op = NVVM_MMA_OPS.all_wmma_ops in {
+          foreach b1op = NVVM_MMA_B1OPS<op>.ret in {
             if NVVM_WMMA_SUPPORTED<op, layout_a, layout_b, satf, rnd>.ret then {
-            def WMMA_NAME<layout_a, layout_b, satf, rnd,
+              def WMMA_NAME<layout_a, layout_b, satf, rnd, b1op,
                                 op[0], op[1], op[2], op[3]>.record
-              : NVVM_WMMA_MMA<layout_a, layout_b, satf, rnd,
+                : NVVM_WMMA_MMA<layout_a, layout_b, satf, rnd, b1op,
                                 op[0], op[1], op[2], op[3]>;
             }
+          } // b1op
         } // op
       } // rnd
     } // satf
   } // layout_b
 } // layout_a
 
 // MMA
-class NVVM_MMA<string ALayout, string BLayout, int Satfinite,
+class NVVM_MMA<string ALayout, string BLayout, int Satfinite, string b1op,
                WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D>
   : Intrinsic<D.regs,
               !listconcat(A.regs, B.regs, C.regs),
               [IntrNoMem],
-              MMA_NAME<ALayout, BLayout, Satfinite, A, B, C, D>.llvm>;
+              MMA_NAME<ALayout, BLayout, Satfinite, b1op, A, B, C, D>.llvm>;
 
 foreach layout_a = ["row", "col"] in {
   foreach layout_b = ["row", "col"] in {
     foreach satf = [0, 1] in {
       foreach op = NVVM_MMA_OPS.all_mma_ops in {
+        foreach b1op = NVVM_MMA_B1OPS<op>.ret in {
           if NVVM_MMA_SUPPORTED<op, layout_a, layout_b, satf>.ret then {
-          def MMA_NAME<layout_a, layout_b, satf, op[0], op[1], op[2], op[3]>.record
-            : NVVM_MMA<layout_a, layout_b, satf, op[0], op[1], op[2], op[3]>;
+            def MMA_NAME<layout_a, layout_b, satf, b1op, op[0], op[1], op[2], op[3]>.record
+              : NVVM_MMA<layout_a, layout_b, satf, b1op, op[0], op[1], op[2], op[3]>;
           }
+        } // b1op
       } // op
     } // satf
   } // layout_b
 } // layout_a
 
 } // let TargetPrefix = "nvvm"

llvm/lib/Target/NVPTX/NVPTXInstrInfo.td

 
 def hasPTX31 : Predicate<"Subtarget->getPTXVersion() >= 31">;
 def hasPTX60 : Predicate<"Subtarget->getPTXVersion() >= 60">;
 def hasPTX61 : Predicate<"Subtarget->getPTXVersion() >= 61">;
 def hasPTX63 : Predicate<"Subtarget->getPTXVersion() >= 63">;
 def hasPTX64 : Predicate<"Subtarget->getPTXVersion() >= 64">;
 def hasPTX65 : Predicate<"Subtarget->getPTXVersion() >= 65">;
 def hasPTX70 : Predicate<"Subtarget->getPTXVersion() >= 70">;
+def hasPTX71 : Predicate<"Subtarget->getPTXVersion() >= 71">;
 
 def hasSM30 : Predicate<"Subtarget->getSmVersion() >= 30">;
 def hasSM70 : Predicate<"Subtarget->getSmVersion() >= 70">;
 def hasSM72 : Predicate<"Subtarget->getSmVersion() >= 72">;
 def hasSM75 : Predicate<"Subtarget->getSmVersion() >= 75">;
 def hasSM80 : Predicate<"Subtarget->getSmVersion() >= 80">;
 
 // non-sync shfl instructions are not available on sm_70+ in PTX6.4+

llvm/lib/Target/NVPTX/NVPTXIntrinsics.td

foreach stride = [false, true] in {
             if NVVM_WMMA_LDST_SUPPORTED<frag, layout>.ret then
               def : WMMA_STORE_D<WMMA_REGINFO<frag, "store">, layout, space, stride, addr>;
         } // addr
       } // space
     } // stride
   } // layout
 } // defset
 
+// B1 instruction variants need extra constraints.
+class MMA_OP_PREDICATES<WMMA_REGINFO FragA, string b1op> {
+  string Op = b1op;
+  WMMA_REGINFO Frag = FragA;
+  list<Predicate> ret = !listconcat(
+    FragA.Predicates,
+    !if(!eq(b1op, ".and.popc"), [hasSM80,hasPTX71],[])
+  );
+}
 // WMMA.MMA
 class WMMA_MMA<WMMA_REGINFO FragA, WMMA_REGINFO FragB,
                WMMA_REGINFO FragC, WMMA_REGINFO FragD,
-               string ALayout, string BLayout, int Satfinite, string rnd>
-  : WMMA_INSTR<WMMA_NAME<ALayout, BLayout, Satfinite, rnd, FragA, FragB, FragC, FragD>.record,
+               string ALayout, string BLayout, int Satfinite, string rnd, string b1op>
+  : WMMA_INSTR<WMMA_NAME<ALayout, BLayout, Satfinite, rnd, b1op, FragA, FragB, FragC, FragD>.record,
                          [FragA.Ins, FragB.Ins, FragC.Ins]>,
     // Requires does not seem to have effect on Instruction w/o Patterns.
     // We set it here anyways and propagate to the Pat<> we construct below.
-    Requires<FragA.Predicates> {
+    Requires<MMA_OP_PREDICATES<FragA, b1op>.ret> {
   let OutOperandList = FragD.Outs;
   let InOperandList  = !con(Args, (ins MmaCode:$ptx));
   string TypeList = !cond(
     !eq(FragA.ptx_elt_type, "f16") : "." # FragD.ptx_elt_type
                                      # "." # FragC.ptx_elt_type,
     1: "." # FragD.ptx_elt_type
        # "." # FragA.ptx_elt_type
        # "." # FragB.ptx_elt_type
        # "." # FragC.ptx_elt_type,
   );
   let AsmString = "wmma.mma"
-                  # !if(!eq(FragA.ptx_elt_type, "b1"), ".xor.popc", "")
+                  # b1op
                   # ".sync"
                   # "${ptx:aligned}"
                   # "." # ALayout
                   # "." # BLayout
                   # "." # FragA.geom
                   # !if(!ne(rnd, ""), !strconcat(".", rnd), "")
                   # TypeList
                   # !if(Satfinite, ".satfinite", "") # "\n\t\t"
                   # FragD.regstring # ",\n\t\t"
                   # FragA.regstring # ",\n\t\t"
                   # FragB.regstring # ",\n\t\t"
                   # FragC.regstring # ";";
 }
 
 defset list<WMMA_INSTR> WMMAs  = {
   foreach layout_a = ["row", "col"] in {
     foreach layout_b = ["row", "col"] in {
       foreach satf = [0, 1] in {
         foreach rnd = ["", "rn", "rz", "rm", "rp"] in {
           foreach op = NVVM_MMA_OPS.all_wmma_ops in {
+            foreach b1op = NVVM_MMA_B1OPS<op>.ret in {
               if NVVM_WMMA_SUPPORTED<op, layout_a, layout_b, satf, rnd>.ret then {
                 def : WMMA_MMA<WMMA_REGINFO<op[0], "wmma.mma">,
                               WMMA_REGINFO<op[1], "wmma.mma">,
                               WMMA_REGINFO<op[2], "wmma.mma">,
                               WMMA_REGINFO<op[3], "wmma.mma">,
-                            layout_a, layout_b, satf, rnd>;
+                              layout_a, layout_b, satf, rnd, b1op>;
               }
+            } // b1op
           } // op
         } // rnd
       } // satf
     } // layout_b
   } // layout_a
 } // defset
 
 // MMA
 class MMA<WMMA_REGINFO FragA, WMMA_REGINFO FragB,
                WMMA_REGINFO FragC, WMMA_REGINFO FragD,
-               string ALayout, string BLayout, int Satfinite>
-  : WMMA_INSTR<MMA_NAME<ALayout, BLayout, Satfinite, FragA, FragB, FragC, FragD>.record,
+               string ALayout, string BLayout, int Satfinite, string b1op>
+  : WMMA_INSTR<MMA_NAME<ALayout, BLayout, Satfinite, b1op, FragA, FragB, FragC, FragD>.record,
                         [FragA.Ins, FragB.Ins, FragC.Ins]>,
     // Requires does not seem to have effect on Instruction w/o Patterns.
     // We set it here anyways and propagate to the Pat<> we construct below.
-    Requires<FragA.Predicates> {
+  Requires<MMA_OP_PREDICATES<FragA, b1op>.ret> {
   let OutOperandList = FragD.Outs;
   let InOperandList  = !con(Args, (ins MmaCode:$ptx));
   string TypeList = "." # FragD.ptx_elt_type
                     # "." # FragA.ptx_elt_type
                     # "." # FragB.ptx_elt_type
                     # "." # FragC.ptx_elt_type;
   let AsmString = "mma.sync.aligned."
                   # FragA.geom
                   # "." # ALayout
                   # "." # BLayout
                   # !if(Satfinite, ".satfinite", "")
                   # TypeList
-                  # !if(!eq(FragA.ptx_elt_type, "b1"), ".xor.popc", "") # "\n\t\t"
+                  # b1op # "\n\t\t"
                   # FragD.regstring # ",\n\t\t"
                   # FragA.regstring # ",\n\t\t"
                   # FragB.regstring # ",\n\t\t"
                   # FragC.regstring # ";";
 }
 
 defset list<WMMA_INSTR> MMAs  = {
   foreach layout_a = ["row", "col"] in {
     foreach layout_b = ["row", "col"] in {
       foreach satf = [0, 1] in {
         foreach op = NVVM_MMA_OPS.all_mma_ops in {
+          foreach b1op = NVVM_MMA_B1OPS<op>.ret in {
             if NVVM_MMA_SUPPORTED<op, layout_a, layout_b, satf>.ret then {
               def : MMA<WMMA_REGINFO<op[0], "mma">,
                         WMMA_REGINFO<op[1], "mma">,
                         WMMA_REGINFO<op[2], "mma">,
                         WMMA_REGINFO<op[3], "mma">,
-                      layout_a, layout_b, satf>;
+                        layout_a, layout_b, satf, b1op>;
             }
+          } // b1op
         } // op
       } // satf
     } // layout_b
   } // layout_a
 } // defset
 
 
 // Constructing non-flat DAGs is still a pain. I can't !subst a dag node with a

llvm/test/CodeGen/NVPTX/wmma.py

 # RUN: %python %s --ptx=65 --gpu-arch=75 > %t-ptx65-sm_75.ll
 # RUN: FileCheck %t-ptx65-sm_75.ll < %t-ptx65-sm_75.ll \
 # RUN:           --check-prefixes=INTRINSICS,M16N16,EXTGEOM,INT,SUBINT,MMA,PTX65MMA
 # RUN: FileCheck %t-ptx65-sm_75.ll < %t-ptx65-sm_75.ll \
 # RUN:           --check-prefixes=INTRINSICS
 # RUN: llc < %t-ptx65-sm_75.ll -march=nvptx64 -mcpu=sm_75 -mattr=+ptx65 \
 # RUN:           | FileCheck %t-ptx65-sm_75.ll
 
-# Check all variants of instructions supported by PTX70 on SM80+
-# RUN: %python %s --ptx=70 --gpu-arch=80 > %t-ptx70-sm_80.ll
-# RUN: FileCheck %t-ptx70-sm_80.ll < %t-ptx70-sm_80.ll \
-# RUN:           --check-prefixes=INTRINSICS,M16N16,EXTGEOM,INT,SUBINT,MMA,ALTFLOAT,DOUBLE,PTX65MMA,PTX70MMA
-# RUN: FileCheck %t-ptx70-sm_80.ll < %t-ptx70-sm_80.ll \
+# Check all variants of instructions supported by PTX71 on SM80+
+# RUN: %python %s --ptx=71 --gpu-arch=80 > %t-ptx71-sm_80.ll
+# RUN: FileCheck %t-ptx71-sm_80.ll < %t-ptx71-sm_80.ll \
+# RUN:           --check-prefixes=INTRINSICS,M16N16,EXTGEOM,INT,SUBINT,MMA,ALTFLOAT,DOUBLE,PTX65MMA,PTX71MMA
+# RUN: FileCheck %t-ptx71-sm_80.ll < %t-ptx71-sm_80.ll \
 # RUN:           --check-prefixes=INTRINSICS
-# RUN: llc < %t-ptx70-sm_80.ll -march=nvptx64 -mcpu=sm_80 -mattr=+ptx70 \
-# RUN:           | FileCheck %t-ptx70-sm_80.ll
+# RUN: llc < %t-ptx71-sm_80.ll -march=nvptx64 -mcpu=sm_80 -mattr=+ptx71 \
+# RUN:           | FileCheck %t-ptx71-sm_80.ll
 
 from __future__ import print_function
 
 import argparse
 from itertools import product
 from string import Template
 
 class MMAType:
"""
   test_params["check_c"] = check_pattern(op.c)
   test_params["check_d"] = check_pattern(op.d)
   args = ",\n        ".join(make_wmma_slice_args(frag)
                             for frag in (op.a, op.b, op.c))
   test_params["args"] = args
   print(Template(mma_template).substitute(test_params))
   return (test_params["intrinsic"], test_params["instruction"])
 
+def get_b1_ops(ptx_type):
+  if ptx_type != "b1":
+    return [""]
+  if ptx_version >= 71:
+    return [".xor.popc", ".and.popc"]
+  return [".xor.popc"]
+
 def gen_wmma_mma_tests():
-  wmma_intrinsic_template = "llvm.nvvm.wmma.${geom}.mma.${alayout}.${blayout}${rnd}.${intrinsic_signature}${satf}"
-  wmma_instruction_template = "wmma.mma${mma_variant}.sync${aligned}.${alayout}.${blayout}.${geom}${rnd}.${ptx_signature}${satf}"
+  wmma_intrinsic_template = "llvm.nvvm.wmma.${geom}.mma${b1op}.${alayout}.${blayout}${rnd}.${intrinsic_signature}${satf}"
+  wmma_instruction_template = "wmma.mma${b1op}.sync${aligned}.${alayout}.${blayout}.${geom}${rnd}.${ptx_signature}${satf}"
 
   generated_items=[]
 
   for op, alayout, blayout, rnd, satf in product(
       get_wmma_ops(),
       ["row","col"],
       ["row","col"],
       [".rn", ".rz", ".rm", ".rp", ""],
       [".satfinite", ""]):
 
     if not is_wmma_variant_supported(op, alayout, blayout, rnd, satf):
       continue
 
+    for b1op in get_b1_ops(op.a.mma_type.ptx_type):
       params = {
           "aligned" : ".aligned" if ptx_version >= 63 else "",
           "alayout" : alayout,
           "blayout" : blayout,
           "intrinsic_signature" : wmma_signature(op),
           "ptx_signature" : wmma_ptx_signature(op),
           "satf"  : satf,
           "rnd"   : rnd,
           "geom"  : op.a.geom,
-        "mma_variant" : ".xor.popc" if op.a.mma_type.ptx_type == "b1" else "",
+          "b1op"  : b1op
       }
 
       intrinsic_template = wmma_intrinsic_template
       instruction_template = wmma_instruction_template
 
       generated_items.append(common_mma_test_gen(params, op,
                                                  intrinsic_template, instruction_template))
 
   return generated_items
 
 def gen_mma_tests():
-  mma_intrinsic_template = "llvm.nvvm.mma.${geom}.${alayout}.${blayout}${satf}.${intrinsic_signature}"
-  mma_instruction_template = "mma.sync${aligned}.${geom}.${alayout}.${blayout}${satf}.${ptx_signature}${mma_variant}"
+  mma_intrinsic_template = "llvm.nvvm.mma${b1op}.${geom}.${alayout}.${blayout}${satf}.${intrinsic_signature}"
+  mma_instruction_template = "mma.sync${aligned}.${geom}.${alayout}.${blayout}${satf}.${ptx_signature}${b1op}"
 
   generated_items=[]
 
   for op, alayout, blayout, satf in product(
       get_mma_ops(),
       ["row","col"],
       ["row","col"],
       [".satfinite", ""]):
 
     if not is_mma_variant_supported(op, alayout, blayout, satf):
       continue
 
+    for b1op in get_b1_ops(op.a.mma_type.ptx_type):
       params = {
           "aligned" : ".aligned" if ptx_version >= 63 else "",
           "alayout" : alayout,
           "blayout" : blayout,
           "intrinsic_signature" : mma_signature(op),
           "ptx_signature" : mma_ptx_signature(op),
           "satf"  : satf,
           "geom"  : op.a.geom,
-        "mma_variant" : ".xor.popc" if op.a.mma_type.ptx_type == "b1" else "",
+          "b1op"  : b1op
       }
 
       intrinsic_template = mma_intrinsic_template
       instruction_template = mma_instruction_template
 
       generated_items.append(common_mma_test_gen(params, op,
         intrinsic_template, instruction_template))
 
   return generated_items
 
 # Append complete list of intrinsics and instructions we've generated tests for.
 # Generate set of checks to verify that that we did generate sensible set of
 # tests for the given combination of PTX and SM variants.
 #
 def gen_check_unsupported_ops(items):

 ; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.s8
 ; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.u8
 ; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.u8.s8
 ; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.u4
 ; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.s4
 ; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.u4
 ; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.s4
 
-; PTX70MMA-DAG: mma.m8n8k4.row.col.f64
-; PTX70MMA-DAG: mma.m16n8k4.row.col.tf32
-; PTX70MMA-DAG: mma.m16n8k8.row.col.tf32
-; PTX70MMA-DAG: mma.m16n8k16.row.col.bf16
-; PTX70MMA-DAG: mma.m16n8k8.row.col.bf16
-; PTX70MMA-DAG: mma.m16n8k16.row.col.f16.f16
-; PTX70MMA-DAG: mma.m16n8k16.row.col.f32.f32
-; PTX70MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.u8
-; PTX70MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.s8
-; PTX70MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.u8
-; PTX70MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.s8
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.u8
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.s8
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.u8
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.s8
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.u4
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.s4
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.u4
-; PTX70MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.s4
-; PTX70MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.u4
-; PTX70MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.s4
-; PTX70MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.u4
-; PTX70MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.s4
-; PTX70MMA-DAG: mma.m8n8k128.row.col.b1
-; PTX70MMA-DAG: mma.m16n8k128.row.col.b1
-; PTX70MMA-DAG: mma.m16n8k256.row.col.b1
+; PTX71MMA-DAG: mma.m8n8k4.row.col.f64
+; PTX71MMA-DAG: mma.m16n8k4.row.col.tf32
+; PTX71MMA-DAG: mma.m16n8k8.row.col.tf32
+; PTX71MMA-DAG: mma.m16n8k16.row.col.bf16
+; PTX71MMA-DAG: mma.m16n8k8.row.col.bf16
+; PTX71MMA-DAG: mma.m16n8k16.row.col.f16.f16
+; PTX71MMA-DAG: mma.m16n8k16.row.col.f32.f32
+; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.u8
+; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.s8
+; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.u8
+; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.s8
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.u8
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.s8
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.u8
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.s8
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.u4
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.s4
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.u4
+; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.s4
+; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.u4
+; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.s4
+; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.u4
+; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.s4
+; PTX71MMA-DAG: mma.and.popc.m8n8k128.row.col.b1
+; PTX71MMA-DAG: mma.xor.popc.m8n8k128.row.col.b1
+; PTX71MMA-DAG: mma.and.popc.m16n8k128.row.col.b1
+; PTX71MMA-DAG: mma.xor.popc.m16n8k128.row.col.b1
+; PTX71MMA-DAG: mma.and.popc.m16n8k256.row.col.b1
+; PTX71MMA-DAG: mma.xor.popc.m16n8k256.row.col.b1
 ;
 
 """)
 
   print("; INTRINSICS_LIST_BEGIN")
   for intrinsic, instruction in sorted(items):
     print("; ", intrinsic, " -> ", instruction,"")
   print("; INTRINSICS_LIST_END")