gps/tools/local/circt/include/mlir/Dialect/LLVMIR/NVVMOps.td

//===-- NVVMOps.td - NVVM IR dialect op definition file ----*- tablegen -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This is the NVVM IR operation definition file.
//
//===----------------------------------------------------------------------===//

#ifndef NVVMIR_OPS
#define NVVMIR_OPS

include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
include "mlir/Dialect/LLVMIR/LLVMTypes.td"
include "mlir/Dialect/LLVMIR/NVVMDialect.td"
include "mlir/Dialect/LLVMIR/NVVMEnums.td"
include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.td"
include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
include "mlir/IR/CommonAttrConstraints.td"
include "mlir/IR/EnumAttr.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

def LLVM_PointerGeneric : LLVM_PointerInAddressSpace<0>;
def LLVM_PointerGlobal : LLVM_PointerInAddressSpace<1>;
def LLVM_PointerShared : LLVM_PointerInAddressSpace<3>;
def LLVM_PointerConst : LLVM_PointerInAddressSpace<4>;
def LLVM_PointerLocal : LLVM_PointerInAddressSpace<5>;
def LLVM_PointerTensor : LLVM_PointerInAddressSpace<6>;
def LLVM_PointerSharedCluster : LLVM_PointerInAddressSpace<7>;

//===----------------------------------------------------------------------===//
// NVVM op definitions
//===----------------------------------------------------------------------===//

class NVVM_Op<string mnemonic, list<Trait> traits = []> :
  LLVM_OpBase<NVVM_Dialect, mnemonic, traits> {
}

/// Base class that defines BasicPtxBuilderOpInterface. 
class NVVM_PTXBuilder_Op<string mnemonic, list<Trait> extraTraits = []> :
  LLVM_OpBase<NVVM_Dialect, mnemonic,
              !listconcat([DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>],
                          extraTraits)> {
}

//===----------------------------------------------------------------------===//
// NVVM attribute definitions
//===----------------------------------------------------------------------===//

class NVVM_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
    : AttrDef<NVVM_Dialect, attrName, traits> {
  let mnemonic = attrMnemonic;
}

// Cache Eviction Priority enum definitions
def EvictNormal : I32EnumCase<"EvictNormal", 0, "evict_normal">;
def EvictFirst : I32EnumCase<"EvictFirst", 1, "evict_first">;
def EvictLast : I32EnumCase<"EvictLast", 2, "evict_last">;
def EvictUnchanged : I32EnumCase<"EvictUnchanged", 3, "evict_unchanged">;
def NoAllocate : I32EnumCase<"NoAllocate", 4, "no_allocate">;

def CacheEvictionPriority : I32Enum<"CacheEvictionPriority",
                                    "NVVM Cache Eviction Priority",
                                    [EvictNormal, EvictFirst, EvictLast,   
                                     EvictUnchanged, NoAllocate]> {
  let cppNamespace = "::mlir::NVVM";
}

def CacheEvictionPriorityAttr : EnumAttr<NVVM_Dialect, CacheEvictionPriority, 
                                         "cache_eviction_priority"> {
  let assemblyFormat = "$value";
}

// Memory Space enum definitions
/// Generic memory space identifier.
def MemSpaceGeneric : I32EnumCase<"Generic", 0, "generic">;
/// Global memory space identifier.
def MemSpaceGlobal : I32EnumCase<"Global", 1, "global">;
/// Shared memory space identifier.
def MemSpaceShared : I32EnumCase<"Shared", 3, "shared">;
/// Constant memory space identifier.
def MemSpaceConstant : I32EnumCase<"Constant", 4, "constant">;
/// Local memory space identifier.
def MemSpaceLocal : I32EnumCase<"Local", 5, "local">;
/// Tensor memory space identifier.
/// Tensor memory is available only in arch-accelerated
/// variants from sm100 onwards.
def MemSpaceTensor : I32EnumCase<"Tensor", 6, "tensor">;
/// Distributed shared memory space identifier.
/// Distributed shared memory is available only in sm90+.
def MemSpaceSharedCluster : I32EnumCase<"SharedCluster", 7, "shared_cluster">;

def NVVMMemorySpace : I32Enum<"NVVMMemorySpace", "NVVM Memory Space",
                  [MemSpaceGeneric, MemSpaceGlobal, MemSpaceShared,   
                   MemSpaceConstant, MemSpaceLocal, MemSpaceTensor, 
                   MemSpaceSharedCluster]> {
  let cppNamespace = "::mlir::NVVM";
}

def NVVMMemorySpaceAttr : 
  EnumAttr<NVVM_Dialect, NVVMMemorySpace, "memory_space", [
    DeclareAttrInterfaceMethods<LLVM_LLVMAddrSpaceAttrInterface>,
    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
  ]> {
  let assemblyFormat = "`<` $value `>`";
}

// Attrs describing the scope of the Memory Operation
def MemScopeKindCTA      : I32EnumAttrCase<"CTA", 0, "cta">;
def MemScopeKindCluster  : I32EnumAttrCase<"CLUSTER", 1, "cluster">;
def MemScopeKindGPU      : I32EnumAttrCase<"GPU", 2, "gpu">;
def MemScopeKindSYS      : I32EnumAttrCase<"SYS", 3, "sys">;

def MemScopeKind : I32EnumAttr<"MemScopeKind", "NVVM Memory Scope kind",
  [MemScopeKindCTA, MemScopeKindCluster, MemScopeKindGPU, MemScopeKindSYS]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MemScopeKindAttr : EnumAttr<NVVM_Dialect, MemScopeKind, "mem_scope"> {
  let assemblyFormat = "`<` $value `>`";
}

// Attrs to disambiguate the cta or cluster space within shared memory
def SharedSpaceCTA : I32EnumAttrCase<"shared_cta", 0, "cta">;
def SharedSpaceCluster   : I32EnumAttrCase<"shared_cluster", 1, "cluster">;
def SharedSpace : I32EnumAttr<"SharedSpace", "Shared memory space",
  [SharedSpaceCTA, SharedSpaceCluster]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def SharedSpaceAttr : EnumAttr<NVVM_Dialect, SharedSpace, "shared_space"> {
  let assemblyFormat = "`<` $value `>`";
}

// Attrs describing the Memory Ordering Semantics
def MemOrderKindWeak     : I32EnumAttrCase<"WEAK", 0, "weak">;
def MemOrderKindRelaxed  : I32EnumAttrCase<"RELAXED", 1, "relaxed">;
def MemOrderKindAcquire  : I32EnumAttrCase<"ACQUIRE", 2, "acquire">;
def MemOrderKindRelease  : I32EnumAttrCase<"RELEASE", 3, "release">;
def MemOrderKindAcqRel   : I32EnumAttrCase<"ACQ_REL", 4, "acq_rel">;
def MemOrderKindSC       : I32EnumAttrCase<"SC", 5, "sc">;
def MemOrderKindMMIO     : I32EnumAttrCase<"MMIO", 6, "mmio">;
def MemOrderKindVolatile : I32EnumAttrCase<"VOLATILE", 7, "volatile">;

def MemOrderKind : I32EnumAttr<"MemOrderKind", "NVVM Memory Ordering kind",
  [MemOrderKindWeak, MemOrderKindRelaxed, MemOrderKindAcquire,
    MemOrderKindRelease, MemOrderKindAcqRel, MemOrderKindSC,
    MemOrderKindMMIO, MemOrderKindVolatile]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MemOrderKindAttr : EnumAttr<NVVM_Dialect, MemOrderKind, "mem_order"> {
  let assemblyFormat = "`<` $value `>`";
}

// Attrs for supported Reduction operations
def ReductionKindNone : I32EnumAttrCase<"NONE", 0, "none">;
def ReductionKindAdd  : I32EnumAttrCase<"ADD", 1, "add">;
def ReductionKindAnd  : I32EnumAttrCase<"AND", 2, "and">;
def ReductionKindMax  : I32EnumAttrCase<"MAX", 3, "max">;
def ReductionKindMin  : I32EnumAttrCase<"MIN", 4, "min">;
def ReductionKindOr   : I32EnumAttrCase<"OR", 5, "or">;
def ReductionKindUmax : I32EnumAttrCase<"UMAX", 6, "umax">;
def ReductionKindUmin : I32EnumAttrCase<"UMIN", 7, "umin">;
def ReductionKindXor  : I32EnumAttrCase<"XOR", 8, "xor">;
def ReductionKindFmin : I32EnumAttrCase<"FMIN", 9, "fmin">;
def ReductionKindFmax : I32EnumAttrCase<"FMAX", 10, "fmax">;

/// Enum attribute of the different kinds.
def ReductionKind : I32EnumAttr<"ReductionKind", "NVVM Reduction Kind attribute",
  [ReductionKindAdd, ReductionKindAnd, ReductionKindMax, ReductionKindMin,
   ReductionKindOr, ReductionKindUmax, ReductionKindUmin, ReductionKindXor,
   ReductionKindFmin, ReductionKindFmax]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}

def ReductionKindAttr : EnumAttr<NVVM_Dialect, ReductionKind, "reduction_kind">;

//===----------------------------------------------------------------------===//
// NVVM intrinsic operations
//===----------------------------------------------------------------------===//

class NVVM_IntrOp<string mnem, list<Trait> traits = [],
                  int numResults = 0>
  : LLVM_IntrOpBase<NVVM_Dialect, mnem, "nvvm_" # !subst(".", "_", mnem),
                    /*list<int> overloadedResults=*/[],
                    /*list<int> overloadedOperands=*/[],
                    traits, numResults>;

class NVVM_IntrinsicLoweringOp<string mnemonic, list<Trait> traits = []> : 
    NVVM_Op<mnemonic, traits> {
  string cppClass = !subst("NVVM_", "", NAME);
  let extraClassDeclaration = [{
    static NVVM::IDArgPair
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::IRBuilderBase &builder);
  }];
}

class NVVM_VoidIntrinsicOp<string mnemonic, list<Trait> traits = []> :
  NVVM_IntrinsicLoweringOp<mnemonic, traits> {
  let llvmBuilder = [{
    auto [id, args] = NVVM::}] # cppClass # [{::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, builder.getVoidTy(), args);
  }];
}

class NVVM_SingleResultIntrinsicOp<string mnemonic, list<Trait> traits = [], string resultStr = "$res"> :
  NVVM_IntrinsicLoweringOp<mnemonic, traits> {
  string getFuncCall = "get" # !toupper(!substr(resultStr, 1, 1)) # !substr(resultStr, 2);
  let llvmBuilder = [{
    auto [id, args] = NVVM::}] # cppClass # [{::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
    if (op->getNumResults() > 0)
      }] # resultStr # [{ = createIntrinsicCall(builder, id, moduleTranslation.convertType(op.}] # getFuncCall # [{().getType()), args);
    else
      createIntrinsicCall(builder, id, builder.getVoidTy(), args);
  }];
}

class NVVM_F32UnaryApproxOp<string mnemonic, list<Trait> traits = []> :
    NVVM_SingleResultIntrinsicOp<mnemonic,
        traits # [Pure, SameOperandsAndResultType]> {
  let arguments = (ins F32:$src,
                       DefaultValuedAttr<BoolAttr, "false">:$ftz);
  let results = (outs F32:$res);
  let assemblyFormat = "$src attr-dict `:` type($src)";
}


//===----------------------------------------------------------------------===//
// NVVM special register op definitions
//===----------------------------------------------------------------------===//

class NVVM_PureSpecialRegisterOp<string mnemonic, list<Trait> traits = []> :
  NVVM_IntrOp<mnemonic, !listconcat(traits, [Pure]), 1> {
  let arguments = (ins);
  let results = (outs I32:$res);
  let assemblyFormat = "attr-dict `:` type($res)";
}

class NVVM_SpecialRegisterOp<string mnemonic, Type resultType = I32,
                                              list<Trait> traits = []> :
  NVVM_IntrOp<mnemonic, traits, 1> {
  let arguments = (ins);
  let results = (outs resultType:$res);
  let assemblyFormat = "attr-dict `:` type($res)";
}

class NVVM_PureSpecialRangeableRegisterOp<string mnemonic, list<Trait> traits = []> :
  NVVM_PureSpecialRegisterOp<mnemonic,
    !listconcat(traits,
      [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>])> {
  let arguments = (ins OptionalAttr<LLVM_ConstantRangeAttr>:$range);
  let assemblyFormat = "(`range` $range^)? attr-dict `:` type($res)";
  let llvmBuilder = baseLlvmBuilder # setRangeRetAttrCode # baseLlvmBuilderCoda;
  let mlirBuilder = baseMlirBuilder # importRangeRetAttrCode # baseMlirBuilderCoda;
  let hasVerifier = 1;

  // Backwards-compatibility builder for an unspecified range.
  let builders = [
    OpBuilder<(ins "Type":$resultType), [{
      build($_builder, $_state, resultType, ::mlir::LLVM::ConstantRangeAttr{});
    }]>
  ];

  // Define this method for the InferIntRangeInterface.
  let extraClassDefinition = [{
    // Infer the result ranges based on the range attribute.
    void $cppClass::inferResultRanges(
        ArrayRef<::mlir::ConstantIntRanges> argRanges,
        SetIntRangeFn setResultRanges) {
        nvvmInferResultRanges(getOperation(), getResult(), argRanges, setResultRanges);
    }

    // Verify the range attribute satisfies LLVM ConstantRange constructor requirements.
    ::llvm::LogicalResult $cppClass::verify() {
      return verifyConstantRangeAttr(getOperation(), getRange());
    }
  }];

}

//===----------------------------------------------------------------------===//
// Lane, Warp, SM, Grid index and range
def NVVM_LaneIdOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.laneid">;
def NVVM_WarpSizeOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.warpsize">;
def NVVM_WarpIdOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.warpid">;
def NVVM_WarpDimOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nwarpid">;
def NVVM_SmIdOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.smid">;
def NVVM_SmDimOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nsmid">;
def NVVM_GridIdOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.gridid">;

//===----------------------------------------------------------------------===//
// Lane Mask Comparison Ops
def NVVM_LaneMaskEqOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.eq">;
def NVVM_LaneMaskLeOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.le">;
def NVVM_LaneMaskLtOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.lt">;
def NVVM_LaneMaskGeOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.ge">;
def NVVM_LaneMaskGtOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.gt">;

//===----------------------------------------------------------------------===//
// Thread index and range
def NVVM_ThreadIdXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.tid.x">;
def NVVM_ThreadIdYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.tid.y">;
def NVVM_ThreadIdZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.tid.z">;
def NVVM_BlockDimXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ntid.x">;
def NVVM_BlockDimYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ntid.y">;
def NVVM_BlockDimZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ntid.z">;

//===----------------------------------------------------------------------===//
// Block index and range
def NVVM_BlockIdXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ctaid.x">;
def NVVM_BlockIdYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ctaid.y">;
def NVVM_BlockIdZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.ctaid.z">;
def NVVM_GridDimXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nctaid.x">;
def NVVM_GridDimYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nctaid.y">;
def NVVM_GridDimZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nctaid.z">;

//===----------------------------------------------------------------------===//
// CTA Cluster index and range
def NVVM_ClusterIdXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.clusterid.x", [NVVMRequiresSM<90>]>;
def NVVM_ClusterIdYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.clusterid.y">;
def NVVM_ClusterIdZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.clusterid.z">;
def NVVM_ClusterDimXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nclusterid.x">;
def NVVM_ClusterDimYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nclusterid.y">;
def NVVM_ClusterDimZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.nclusterid.z">;


//===----------------------------------------------------------------------===//
// CTA index and range within Cluster
def NVVM_BlockInClusterIdXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.ctaid.x", [NVVMRequiresSM<90>]>;
def NVVM_BlockInClusterIdYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.ctaid.y", [NVVMRequiresSM<90>]>;
def NVVM_BlockInClusterIdZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.ctaid.z", [NVVMRequiresSM<90>]>;
def NVVM_ClusterDimBlocksXOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.nctaid.x", [NVVMRequiresSM<90>]>;
def NVVM_ClusterDimBlocksYOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.nctaid.y", [NVVMRequiresSM<90>]>;
def NVVM_ClusterDimBlocksZOp : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.nctaid.z">;

//===----------------------------------------------------------------------===//
// CTA index and across Cluster dimensions
def NVVM_ClusterId : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.ctarank", [NVVMRequiresSM<90>]>;
def NVVM_ClusterDim : NVVM_PureSpecialRangeableRegisterOp<"read.ptx.sreg.cluster.nctarank">;

//===----------------------------------------------------------------------===//
// Various configurations of Shared memory sizes
def NVVM_TotalSmemSize   : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.total.smem.size">;
def NVVM_DynamicSmemSize : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.dynamic.smem.size">;
def NVVM_AggrSmemSize    : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.aggr.smem.size", [NVVMRequiresSM<90>]>;

//===----------------------------------------------------------------------===//
// Clock registers
def NVVM_ClockOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.clock">;
def NVVM_Clock64Op : NVVM_SpecialRegisterOp<"read.ptx.sreg.clock64", I64>;
def NVVM_GlobalTimerOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.globaltimer", I64>;
def NVVM_GlobalTimerLoOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.globaltimer.lo">;

//===----------------------------------------------------------------------===//
// envreg registers
foreach index = !range(0, 32) in {
  def NVVM_EnvReg # index # Op : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.envreg" # index>;
}

//===----------------------------------------------------------------------===//
// Inline PTX op definition
//===----------------------------------------------------------------------===//

def NVVM_InlinePtxOp : NVVM_Op<"inline_ptx", 
  [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>, 
    AttrSizedOperandSegments]>
{
  let summary = "Inline PTX Op";
  let description = [{This op allows using PTX directly within the NVVM 
    dialect, while greatly simplifying llvm.inline_asm generation. It 
    automatically handles register size selection and sets the correct 
    read/write access for each operand. The operation leverages the 
    `BasicPtxBuilderInterface` to abstract away low-level details of 
    PTX assembly formatting.

    The `predicate` attribute is used to specify a predicate for the 
    PTX instruction.

    Example 1: Read-only Parameters
    ```mlir
    nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count) : !llvm.ptr, i32

    // Lowers to:
    llvm.inline_asm has_side_effects asm_dialect = att 
      "mbarrier.init.b64 [$0], $1;", "l,r" %arg0, %arg2 : (!llvm.ptr, i32) -> ()
    ```

    Example 2: Read-only and Write-only Parameters
    ```mlir
    %0 = nvvm.inline_ptx "ex2.approx.ftz.f32 $0, $1;" (%input) : f32 -> f32

    // Lowers to:
    %0 = llvm.inline_asm has_side_effects asm_dialect = att 
      "ex2.approx.ftz.f32 $0, $1;", "=f,f" %arg0 : (f32) -> f32
    ```

    Example 3: Predicate Usage
    ```mlir
    nvvm.inline_ptx "mbarrier.init.b64 [$0], $1;" (%barrier_gen, %count), 
      predicate = %pred : !llvm.ptr, i32, i1

    // Lowers to:
    llvm.inline_asm has_side_effects asm_dialect = att 
      "@$2 mbarrier.init.b64 [$0], $1;", "l,r,b" %arg0, %arg2, %arg3 
      : (!llvm.ptr, i32, i1) -> ()
    ```
  }];

  let arguments = (ins Variadic<AnyType>:$readOnlyArgs, 
                       Variadic<AnyType>:$readWriteArgs, 
                       StrAttr:$ptxCode,
                       PtxPredicate:$predicate);
                      
  let results = (outs Variadic<AnyType>:$writeOnlyArgs);
  
  let assemblyFormat = [{
    $ptxCode
    ( `ro` `(` $readOnlyArgs^ `:` type($readOnlyArgs) `)` )?
    ( `rw` `(` $readWriteArgs^ `:` type($readWriteArgs) `)` )?
    (`,` `predicate` `=` $predicate^)? 
    attr-dict
    ( `->` type($writeOnlyArgs)^ )?
  }];
  
  let extraClassDefinition = [{
    std::string $cppClass::getPtx() {
      StringRef ptxInstStr = getPtxCode();
      return std::string(ptxInstStr.data());
    }
  }];

  let extraClassDeclaration = [{
    bool getAsmValues(RewriterBase &, llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM approximate op definitions
//===----------------------------------------------------------------------===//

def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
  let arguments = (ins F32:$arg);
  let results = (outs F32:$res);
  let assemblyFormat = "$arg attr-dict `:` type($res)";
}

//===----------------------------------------------------------------------===//
// NVVM transcendental op definitions
//===----------------------------------------------------------------------===//
// PTX provides only fast approximations for sin/cos/log2/ex2/tanh. Exposing
// these as plain `nvvm.{sin,cos,log2,ex2,tanh}` (without a `.approx` suffix)
// keeps the NVVM dialect self-contained -- users don't need to reach for
// `math.*` ops for something that already has a PTX instruction.

def NVVM_SinOp : NVVM_F32UnaryApproxOp<"sin"> {
  let summary = "Sine (fast approximation)";
  let description = [{
    Computes a fast approximation of the sine of the input value (in radians).
    The `ftz` attribute, when set, flushes subnormal inputs and results to
    sign-preserving zero.

    For more information, see PTX ISA:
    [sin](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sin)
  }];
}

def NVVM_CosOp : NVVM_F32UnaryApproxOp<"cos"> {
  let summary = "Cosine (fast approximation)";
  let description = [{
    Computes a fast approximation of the cosine of the input value (in
    radians). The `ftz` attribute, when set, flushes subnormal inputs
    and results to sign-preserving zero.
  }];
}

def NVVM_Log2Op : NVVM_F32UnaryApproxOp<"log2"> {
  let summary = "Base-2 logarithm (fast approximation)";
  let description = [{
    Computes a fast approximation of the base-2 logarithm of the input
    value. The `ftz` attribute, when set, flushes subnormal inputs and
    results to sign-preserving zero.
  }];
}

def NVVM_Ex2Op : NVVM_F32UnaryApproxOp<"ex2"> {
  let summary = "Base-2 exponential (fast approximation)";
  let description = [{
    Computes a fast approximation of 2 raised to the power of the input
    value. The `ftz` attribute, when set, flushes subnormal inputs and
    results to sign-preserving zero.
  }];
}

//===----------------------------------------------------------------------===//
// NVVM rsqrt op definitions
//===----------------------------------------------------------------------===//

def NVVM_RsqrtOp
    : NVVM_SingleResultIntrinsicOp<"rsqrt", [Pure, SameOperandsAndResultType]> {
  let summary = "Reciprocal square root (fast approximation)";
  let description = [{
    Computes an approximation of the reciprocal of the square root of the
    input value: `d = 1 / sqrt(a)`. Supports both f32 and f64. The maximum
    relative error for the f32 form over the entire positive finite range
    is 2^-22.9.

    The `ftz` attribute, when set, flushes subnormal inputs and results to
    sign-preserving zero. For f64 inputs, `ftz=true` selects a coarser
    approximation that uses only the upper 32 bits of the input (the lower
    32 bits of the result are zeroed).

    For more information, see PTX ISA:
    [rsqrt](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-rsqrt)
  }];

  let arguments = (ins AnyTypeOf<[F32, F64]>:$src,
      DefaultValuedAttr<BoolAttr, "false">:$ftz);
  let results = (outs AnyTypeOf<[F32, F64]>:$res);
  let assemblyFormat = "$src attr-dict `:` type($src)";
}

//===----------------------------------------------------------------------===//
// NVVM redux op definitions
//===----------------------------------------------------------------------===//

def NVVM_ReduxOp :
  NVVM_Op<"redux.sync", [NVVMRequiresSM<80>, AllTypesMatch<["res", "val"]>]>,
  Results<(outs AnyTypeOf<[I32, F32]>:$res)>,
  Arguments<(ins AnyTypeOf<[I32, F32]>:$val,
                 ReductionKindAttr:$kind,
                 I32:$mask_and_clamp,
                 DefaultValuedAttr<BoolAttr, "false">:$abs,
                 DefaultValuedAttr<BoolAttr, "false">:$nan)> {
  let summary = "Redux Sync Op";
  let description = [{
    `redux.sync` performs a reduction operation `kind` of the 32 bit source 
    register across all non-exited threads in the membermask.

    The `abs` and `nan` attributes can be used in the case of f32 input type, 
    where the `abs` attribute causes the absolute value of the input to be used 
    in the reduction operation, and the `nan` attribute causes the reduction 
    operation to return NaN if any of the inputs to participating threads are 
    NaN.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync)
  }];
  let hasVerifier = 1;

  string llvmBuilder = [{
      auto intId = getReduxIntrinsicId($_resultType, $kind, $abs, $nan);
      $res = createIntrinsicCall(builder, intId, {$val, $mask_and_clamp});
  }];
  let assemblyFormat = [{
    $kind $val `,` $mask_and_clamp  attr-dict `:` type($val) `->` type($res)
   }];   
}

//===----------------------------------------------------------------------===//
// NVVM nanosleep
//===----------------------------------------------------------------------===//

def NVVM_NanosleepOp : NVVM_Op<"nanosleep">,
  Arguments<(ins I32:$duration)> 
{
  let summary = "Suspends the thread for a specified duration.";

  let description = [{
    The op suspends the thread for a sleep duration approximately close to the 
    delay `$duration`, specified in nanoseconds. 

    The sleep duration is approximated, but guaranteed to be in the 
    interval [0, 2*t]. The maximum sleep duration is 1 millisecond. 
    The implementation may reduce the sleep duration for individual threads 
    within a warp such that all sleeping threads in the warp wake up together.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-nanosleep)
  }];
  
  string llvmBuilder = [{
      createIntrinsicCall(builder, 
                          llvm::Intrinsic::nvvm_nanosleep, {$duration});
  }];
  let assemblyFormat = "attr-dict $duration";
}

//===----------------------------------------------------------------------===//
// NVVM Performance Monitor events
//===----------------------------------------------------------------------===//

def NVVM_PMEventOp : NVVM_VoidIntrinsicOp<"pmevent">,
  Arguments<(ins OptionalAttr<I16Attr>:$maskedEventId, 
                 OptionalAttr<I32Attr>:$eventId)> {
  let summary = "Trigger one or more Performance Monitor events.";

  let description = [{
    Triggers one or more of a fixed number of performance monitor events, with
    event index or mask specified by immediate operand.

    Without `mask` it triggers a single performance monitor event indexed by
    immediate operand a, in the range 0..15.

    With `mask` it triggers one or more of the performance monitor events. Each
    bit in the 16-bit immediate operand controls an event.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-pmevent)
  }];
  
  let assemblyFormat = "attr-dict (`id` `=` $eventId^)? (`mask` `=` $maskedEventId^)?";
  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// NVVM Split arrive/wait barrier
//===----------------------------------------------------------------------===//

/// mbarrier.init instruction with generic pointer type
def NVVM_MBarrierInitOp : NVVM_PTXBuilder_Op<"mbarrier.init">,
  Arguments<(ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
                 I32:$count, PtxPredicate:$predicate)> {
  let summary = "MBarrier Initialization Op";
  let description = [{
    The `nvvm.mbarrier.init` operation initializes an *mbarrier object* at the specified 
    memory location. 

    This operation initializes the *mbarrier object* with the following state:
    - Current phase: 0
    - Expected arrival count: `count`
    - Pending arrival count: `count`  
    - Transaction count (tx-count): 0

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. The `addr`
      must be a pointer to generic or shared::cta memory. When it is generic, the
      underlying address must be within the shared::cta memory space; otherwise
      the behavior is undefined.
    - `count`: Integer specifying the number of threads that will participate in barrier
      synchronization. Must be in the range [1, 2²⁰ - 1].
    - `predicate`: Optional predicate for conditional execution.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init)
  }];
  let assemblyFormat = "$addr `,` $count (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";

  let extraClassDeclaration = [{
    bool hasIntrinsic() { if(getPredicate()) return false; return true; }

    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::MBarrierInitOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_MBarrierInvalOp : NVVM_VoidIntrinsicOp<"mbarrier.inval">,
  Arguments<(ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr)> {
  let summary = "MBarrier Invalidation Operation";
  let description = [{
    The `nvvm.mbarrier.inval` operation invalidates an *mbarrier object* at the 
    specified memory location. 

    This operation marks the *mbarrier object* as invalid, making it safe to repurpose 
    the memory location for other uses or to reinitialize it as a new *mbarrier object*.
    It is undefined behavior if the *mbarrier object* is already invalid.
    
    The operation takes the following operand:
    - `addr`: A pointer to the memory location of the *mbarrier object*. The `addr`
      must be a pointer to generic or shared::cta memory. When it is generic, the
      underlying address must be within the shared::cta memory space; otherwise
      the behavior is undefined.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-inval)
  }];

  let assemblyFormat = "$addr attr-dict `:` type(operands)";
}

def NVVM_MBarrierExpectTxOp : NVVM_VoidIntrinsicOp<"mbarrier.expect_tx"> {
  let summary = "MBarrier expect-tx Operation";
  let description = [{
    The `nvvm.mbarrier.expect_tx` operation increases the transaction count
    of the mbarrier located at `addr` by `txcount` amount. The `scope`
    specifies the set of threads that can directly observe the memory
    synchronizing effect of the `mbarrier.expect_tx` operation. `CTA`
    and `CLUSTER` are the only allowed values for `scope`.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-expect-tx)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    I32:$txcount,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope);

  let assemblyFormat = "$addr `,` $txcount attr-dict `:` type(operands)";

  let hasVerifier = 1;
}

def NVVM_MBarrierCompleteTxOp : NVVM_VoidIntrinsicOp<"mbarrier.complete_tx"> {
  let summary = "MBarrier complete-tx Operation";
  let description = [{
    The `nvvm.mbarrier.complete_tx` operation decrements the transaction
    count of the *mbarrier object* at `addr` by `txcount`. It also signals
    the completion of asynchronous transactions that were tracked by the
    current phase. The `scope` specifies the set of threads that can directly
    observe the memory synchronizing effect of the `mbarrier.complete_tx`
    operation. `CTA` and `CLUSTER` are the only allowed values for `scope`.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-complete-tx)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    I32:$txcount,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope);

  let assemblyFormat = "$addr `,` $txcount attr-dict `:` type(operands)";

  let hasVerifier = 1;
}

def NVVM_MBarrierArriveOp : NVVM_SingleResultIntrinsicOp<"mbarrier.arrive",
    [InferTypeOpAdaptorWithIsCompatible]> {
  let summary = "MBarrier Arrive Operation";
  let description = [{
    The `nvvm.mbarrier.arrive` operation performs an arrive-on operation on the 
    *mbarrier object* at the specified address. Uses the default `.release.cta` semantics. 
    This release pattern establishes memory ordering for operations occurring in program 
    order before this arrive instruction by making operations from the current thread 
    visible to subsequent operations in other threads within the CTA. When other threads 
    perform corresponding acquire operations (like 'mbarrier.test.wait'), they synchronize 
    with this release pattern.

    This operation causes the executing thread to signal its arrival at the barrier.

    - `res`: When the `space` is not shared_cluster, this operation returns an
      opaque 64-bit value capturing the phase of the *mbarrier object* prior to
      the arrive-on operation. The contents of this return value are
      implementation-specific. An *mbarrier object* located in the shared_cluster
      space cannot return a value.

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. The `addr`
      must be a pointer to generic or shared_cta or shared_cluster memory. When it
      is generic, the underlying address must be within the shared_cta memory space;
      otherwise the behavior is undefined.
    - `count`: This specifies the amount by which the pending arrival count is
      decremented. If the `count` argument is not specified, the pending arrival
      count is decremented by 1.
    - `scope`: This specifies the set of threads that directly observe the memory
      synchronizing effect of the `mbarrier.arrive` operation.
    - `space`: This indicates the memory space where the mbarrier object resides.
    - `relaxed`: When set to true, the `arrive` operation has relaxed memory semantics
      and does not provide any ordering or visibility guarantees.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive)
  }];

  let results = (outs Optional<I64>:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    Optional<I32>:$count,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed);

  let assemblyFormat = "$addr (`,` $count^)? attr-dict `:` type($addr) (`->` type($res)^)?";

  let hasVerifier = 1;
}

def NVVM_MBarrierArriveDropOp : NVVM_SingleResultIntrinsicOp<"mbarrier.arrive_drop",
    [InferTypeOpAdaptorWithIsCompatible]> {
  let summary = "MBarrier Arrive-Drop Operation";
  let description = [{
    The `nvvm.mbarrier.arrive_drop` operation decrements the expected arrival
    count of the *mbarrier object* by `count` and then performs an arrive-on
    operation. When `count` is not specified, it defaults to 1. The decrement
    of the expected arrival count applies to all the subsequent phases of the
    *mbarrier object*. The remaining semantics are identical to those of the
    `nvvm.mbarrier.arrive` operation.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive-drop)
  }];

  let results = (outs Optional<I64>:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    Optional<I32>:$count,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed);

  let assemblyFormat = "$addr (`,` $count^)? attr-dict `:` type($addr) (`->` type($res)^)?";

  let hasVerifier = 1;
}

def NVVM_MBarrierArriveNocompleteOp : NVVM_SingleResultIntrinsicOp<"mbarrier.arrive.nocomplete">,
  Results<(outs I64:$res)>,
  Arguments<(ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
                 I32:$count)> {
  let summary = "MBarrier Arrive No-Complete Operation";
  let description = [{
    The `nvvm.mbarrier.arrive.nocomplete` operation performs an arrive-on operation 
    on the *mbarrier object* with the guarantee that it will not cause the barrier to 
    complete its current phase. Uses the default `.release.cta` semantics. This release 
    pattern establishes memory ordering for operations occurring in program order before 
    this arrive instruction by making operations from the current thread visible to 
    subsequent operations in other threads within the CTA. When other threads perform 
    corresponding acquire operations (like 'mbarrier.test.wait'), they synchronize with 
    this release pattern.

    This operation causes the executing thread to signal its arrival at the barrier 
    with a specified count, but ensures that the barrier phase will not complete as 
    a result of this operation. The operation returns an opaque value that 
    captures the phase of the *mbarrier object* prior to the arrive-on operation.

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. The `addr`
      must be a pointer to generic or shared::cta memory. When it is generic, the
      underlying address must be within the shared::cta memory space; otherwise
      the behavior is undefined.
    - `count`: Integer specifying the count argument to the arrive-on operation. 
      Must be in the valid range as specified in the *mbarrier object* contents.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive)
  }];

  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
}

def NVVM_MBarrierArriveDropNocompleteOp : NVVM_SingleResultIntrinsicOp<"mbarrier.arrive_drop.nocomplete">,
  Results<(outs I64:$res)>,
  Arguments<(ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
                 I32:$count)> {
  let summary = "MBarrier Arrive-Drop No-Complete Operation";
  let description = [{
    The `nvvm.mbarrier.arrive_drop.nocomplete` operation decrements the expected
    arrival count of the *mbarrier object* by the amount `count` and then performs
    an arrive-on operation on the *mbarrier object* with the guarantee that it
    will not cause the barrier to complete its current phase.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive-drop)
  }];

  let assemblyFormat = "$addr `,` $count attr-dict `:` type(operands) `->` type($res)";
}

def NVVM_MBarrierArriveExpectTxOp : NVVM_PTXBuilder_Op<"mbarrier.arrive.expect_tx",
    [InferTypeOpAdaptorWithIsCompatible]> {
  let summary = "MBarrier Arrive with Expected Transaction Count";
  let description = [{
    The `nvvm.mbarrier.arrive.expect_tx` operation performs an expect-tx operation 
    followed by an arrive-on operation on the *mbarrier object*. Uses the default 
    `.release.cta` semantics. This release pattern establishes memory ordering for 
    operations occurring in program order before this arrive instruction by making 
    operations from the current thread visible to subsequent operations in other 
    threads within the CTA. When other threads perform corresponding acquire operations 
    (like 'mbarrier.test.wait'), they synchronize with this release pattern.

    This operation first performs an expect-tx operation with the specified transaction
    count, then performs an arrive-on operation with an implicit count of 1. The
    expect-tx operation increases the expect-count of the *mbarrier object* by the
    specified value (i.e. `txcount`), setting the current phase to expect and track
    the completion of additional asynchronous transactions.

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. Uses generic 
      addressing, but the address must still be in the shared memory space.
    - `txcount`: An unsigned integer specifying the expected transaction count 
      for the expect-tx operation. This represents the number of asynchronous transactions 
      expected to complete before the barrier phase completes.
    - `scope`: This specifies the set of threads that directly observe the memory
      synchronizing effect of the `mbarrier.test.wait` operation.
    - `relaxed`: When set to true, the `arrive` operation has relaxed memory semantics
      and does not provide any ordering or visibility guarantees.
    - `predicate`: Optional predicate for conditional execution used only when lowering to
      inline-ptx.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive-drop)
  }];

  let results = (outs Optional<I64>:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    I32:$txcount,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed,
    PtxPredicate:$predicate);

  let assemblyFormat = "$addr `,` $txcount (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands) (`->` type($res)^)?";
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    bool hasIntrinsic() { return !getPredicate(); }

    bool getAsmValues(RewriterBase &rewriter,
      llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues);

    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::MBarrierArriveExpectTxOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);

    if (op.getNumResults() > 0)
      $res = createIntrinsicCall(builder, id, args);
    else
      createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_MBarrierArriveDropExpectTxOp : NVVM_SingleResultIntrinsicOp<"mbarrier.arrive_drop.expect_tx",
    [InferTypeOpAdaptorWithIsCompatible]> {
  let summary = "MBarrier arrive_drop with expected transaction count";
  let description = [{
    The `nvvm.mbarrier.arrive_drop.expect_tx` operation is similar to the
    `nvvm.mbarrier.arrive.expect_tx` operation except that it performs an
    `arrive_drop` operation instead of only an `arrive` operation.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive-drop)
  }];

  let results = (outs Optional<I64>:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared, LLVM_PointerSharedCluster]>:$addr,
    I32:$txcount,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed);

  let assemblyFormat = "$addr `,` $txcount attr-dict `:` type(operands) (`->` type($res)^)?";
  let hasVerifier = 1;
}

def NVVM_MBarrierTryWaitParityOp : NVVM_PTXBuilder_Op<"mbarrier.try_wait.parity">,  
  Arguments<(ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
    I32:$phase, I32:$ticks)> {
  let summary = "MBarrier Potentially-Blocking Try Wait with Phase Parity";
  let description = [{
    The `nvvm.mbarrier.try_wait.parity` operation performs a potentially-blocking 
    test for the completion of a specific phase of an *mbarrier object* using phase 
    parity. It uses the default `.acquire.cta` semantics. This acquire pattern 
    establishes memory ordering for operations occurring in program order after this 
    wait instruction by making operations from other threads in the CTA visible to subsequent 
    operations in the current thread. When this wait completes, it synchronizes with 
    the corresponding release pattern from the `mbarrier.arrive` operation, establishing 
    memory ordering within the CTA.

    This operation waits for the completion of the mbarrier phase indicated by the 
    phase parity. While it uses the underlying PTX `mbarrier.try_wait.parity` 
    instruction, this MLIR operation generates a loop that enforces the test to 
    complete before continuing execution, ensuring the barrier phase is actually 
    completed rather than potentially timing out.

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. Uses generic 
      addressing, but the address must still be in the shared memory space.
    - `phase`: An integer specifying the phase parity (0 or 1). Even phases 
      have parity 0, odd phases have parity 1.
    - `ticks`: An unsigned integer specifying the suspend time hint in 
      nanoseconds. This may be used instead of the system-dependent time limit.

    **Memory ordering guarantees**: When this wait returns true, the following 
    ordering guarantees hold:

    1. All memory accesses (except async operations) requested prior to
       `mbarrier.arrive` having release semantics by participating CTA threads
       are visible to the executing thread.
    2. All `cp.async` operations requested prior to `cp.async.mbarrier.arrive`
       by participating CTA threads are visible to the executing thread.
    3. All `cp.async.bulk` operations using the same *mbarrier object* requested
       prior to `mbarrier.arrive` having release semantics by participating CTA
       threads are visible to the executing thread.
    4. Memory accesses requested after this wait are not visible to memory
       accesses performed prior to `mbarrier.arrive` by other participating
       threads.
    5. No ordering guarantee exists for memory accesses by the same thread
       between `mbarrier.arrive` and this wait.

    **Implementation behavior**:
    This operation generates a PTX loop that repeatedly calls the underlying 
    `mbarrier.try_wait.parity` instruction until the barrier phase completes. 
    Unlike the raw PTX instruction which may return without completion after a 
    timeout, this MLIR operation guarantees completion by continuing to loop until 
    the specified phase is reached.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-try-wait)
  }];  
  let assemblyFormat = "$addr `,` $phase `,` $ticks attr-dict `:` type(operands)";
}

def NVVM_MBarrierTestWaitOp : NVVM_SingleResultIntrinsicOp<"mbarrier.test.wait"> {
  let summary = "MBarrier Non-Blocking Test Wait Operation";
  let description = [{
    The `nvvm.mbarrier.test.wait` operation performs a non-blocking test for the
    completion of a specific phase of an *mbarrier object*. It uses the default
    `.acquire.cta` semantics. This acquire pattern establishes memory ordering for
    operations occurring in program order after this wait instruction by making
    operations from other threads in the CTA visible to subsequent operations in the current
    thread. When this wait completes, it synchronizes with the corresponding release
    pattern from the `mbarrier.arrive` operation, establishing memory ordering within
    the CTA.

    This operation tests whether the mbarrier phase specified by the state operand
    has completed. It is a non-blocking instruction that immediately returns the
    completion status without suspending the executing thread.

    The operation takes the following operands:
    - `addr`: A pointer to the memory location of the *mbarrier object*. Uses generic
      addressing, but the address must still be in the shared memory space.
    - `stateOrPhase`: This argument represents a `state` when it is a 64-bit value
      and represents a `phase` when it is a 32-bit value. The `state` is an opaque
      value returned by a previous `mbarrier.arrive` operation on the same
      *mbarrier object* during the current or immediately preceding phase.
      The `phase` is an integer specifying the phase parity (0 or 1).
      Even phases have parity 0, odd phases have parity 1.
    - `scope`: This specifies the set of threads that directly observe the memory
      synchronizing effect of the `mbarrier.test.wait` operation.
    - `relaxed`: When set to true, the `arrive` operation has relaxed memory semantics
      and does not provide any ordering or visibility guarantees.

    The operation returns a boolean value indicating whether the specified phase 
    has completed:
    - `true`: The immediately preceding phase has completed
    - `false`: The phase is still incomplete (current phase)

    **Memory ordering guarantees**: When this wait returns true, the following 
    ordering guarantees hold:

    1. All memory accesses (except async operations) requested prior to
       `mbarrier.arrive` having release semantics by participating CTA threads
       are visible to the executing thread.
    2. All `cp.async` operations requested prior to `cp.async.mbarrier.arrive`
       by participating CTA threads are visible to the executing thread.
    3. All `cp.async.bulk` operations using the same *mbarrier object* requested
       prior to `mbarrier.arrive` having release semantics by participating CTA
       threads are visible to the executing thread.
    4. Memory accesses requested after this wait are not visible to memory
       accesses performed prior to `mbarrier.arrive` by other participating
       threads.
    5. No ordering guarantee exists for memory accesses by the same thread
       between `mbarrier.arrive` and this wait.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-try-wait)
  }];

  let results = (outs I1:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
    AnyTypeOf<[I64, I32]>:$stateOrPhase,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed);

  let assemblyFormat = "$addr `,` $stateOrPhase attr-dict `:` type(operands) `->` type($res)";
  let hasVerifier = 1;
}

def NVVM_MBarrierTryWaitOp : NVVM_SingleResultIntrinsicOp<"mbarrier.try_wait"> {
  let summary = "MBarrier try wait on state or phase with an optional timelimit";
  let description = [{
    The `nvvm.mbarrier.try_wait` operation checks whether the specified
    *mbarrier object* at `addr` has completed the given phase. Note that
    unlike the `nvvm.mbarrier.test.wait` operation, the try_wait operation
    is a potentially-blocking one. If the phase is not yet complete, the
    calling thread may be suspended. A suspended thread resumes execution
    once the phase completes or when a system-defined timeout occurs.
    Optionally, the `ticks` operand can be used to provide a custom timeout
    (in nanoseconds), overriding the system-defined one. The semantics of
    this operation and its operands are otherwise similar to those of the
    `nvvm.mbarrier.test.wait` Op.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-try-wait)
  }];

  let results = (outs I1:$res);
  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
    AnyTypeOf<[I64, I32]>:$stateOrPhase,
    Optional<I32>:$ticks,
    DefaultValuedAttr<MemScopeKindAttr, "MemScopeKind::CTA">:$scope,
    DefaultValuedAttr<BoolAttr, "false">:$relaxed);

  let assemblyFormat = "$addr `,` $stateOrPhase (`,` $ticks^)? attr-dict `:` type(operands) `->` type($res)";
  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// NVVM synchronization op definitions
//===----------------------------------------------------------------------===//

// Attrs describing the reduction operations for the barrier operation.
def BarrierReductionPopc : I32EnumAttrCase<"POPC", 0, "popc">;
def BarrierReductionAnd : I32EnumAttrCase<"AND", 1, "and">;
def BarrierReductionOr : I32EnumAttrCase<"OR", 2, "or">;

def BarrierReduction
    : I32EnumAttr<"BarrierReduction", "NVVM barrier reduction operation",
                  [BarrierReductionPopc, BarrierReductionAnd,
                   BarrierReductionOr]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def BarrierReductionAttr
    : EnumAttr<NVVM_Dialect, BarrierReduction, "reduction"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_BarrierOp : NVVM_VoidIntrinsicOp<"barrier",
    [AttrSizedOperandSegments]> {
  let summary = "CTA Barrier Synchronization Op";
  let description = [{
    The `nvvm.barrier` operation performs barrier synchronization and communication 
    within a CTA (Cooperative Thread Array). It causes executing threads to wait for 
    all non-exited threads participating in the barrier to arrive.

    The operation takes the following optional operands and attributes:

    - `barrierId`: Specifies a logical barrier resource with value 0 through 15. 
      Each CTA instance has sixteen barriers numbered 0..15. Defaults to 0 if not specified.
    - `numberOfThreads`: Specifies the number of threads participating in the barrier. 
      When specified, the value must be a multiple of the warp size. If not specified, 
      all threads in the CTA participate in the barrier.
    - `aligned`: Selects between the `.aligned` and non-`.aligned` forms of the
      underlying `@llvm.nvvm.barrier.cta.*` intrinsic family. Defaults to true,
      which requires every thread in the CTA to reach this same barrier
      instruction, otherwise the behavior is undefined. Set it to false to emit
      the non-`.aligned` form.

    Reduction variants of the barrier instruction are modeled by the
    `nvvm.barrier.reduction` op.

    The barrier operation guarantees that when the barrier completes, prior memory 
    accesses requested by participating threads are performed relative to all threads 
    participating in the barrier. It also ensures that no new memory access is 
    requested by participating threads before the barrier completes.

    When a barrier completes, the waiting threads are restarted without delay, and
    the barrier is reinitialized so that it can be immediately reused.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
  }];

  let arguments = (ins
      Optional<I32>:$barrierId,
      Optional<I32>:$numberOfThreads,
      DefaultValuedAttr<BoolAttr, "true">:$aligned);

  let assemblyFormat =
      "(`id` `=` $barrierId^)? (`number_of_threads` `=` $numberOfThreads^)? "
      "attr-dict";

  let builders = [OpBuilder<(ins), [{
      return build($_builder, $_state, Value{}, Value{});
    }]>];
}

def NVVM_BarrierReductionOp :
    NVVM_IntrinsicLoweringOp<"barrier.reduction"> {
  let summary = "CTA Barrier Reduction Op";
  let description = [{
    The `nvvm.barrier.reduction` operation performs barrier synchronization with a
    reduction across the per-thread predicates contributed by participating threads
    in a CTA.

    - `barrierId`: Specifies a logical barrier resource with value 0 through 15.
      Optional; defaults to barrier id 0 when not specified.
    - `reductionOp`: The reduction kind (`popc`, `and`, `or`) applied across the
      per-thread predicates.
    - `reductionPredicate`: The per-thread i32 predicate. It is compared against
      zero to form the i1 value fed into the reduction.
    - `aligned`: Selects between the `.aligned` and non-`.aligned` forms of the
      underlying `@llvm.nvvm.barrier.cta.red.*` intrinsic family. Defaults to
      true, which requires every thread in the CTA to reach this same barrier
      instruction, otherwise the behavior is undefined. Set it to false to emit
      the non-`.aligned` form.

    The result is the i32 reduction value computed across all threads
    participating in the barrier.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
  }];

  let arguments = (ins
      Optional<I32>:$barrierId,
      BarrierReductionAttr:$reductionOp,
      I32:$reductionPredicate,
      DefaultValuedAttr<BoolAttr, "true">:$aligned);
  let results = (outs I32:$res);

  let assemblyFormat =
      "qualified($reductionOp) $reductionPredicate "
      "(`id` `=` $barrierId^)? `->` type($res) attr-dict";

  string llvmBuilder = [{
    auto [id, args] = NVVM::BarrierReductionOp::getIntrinsicIDAndArgs(
                        *op, moduleTranslation, builder);
    $res = builder.CreateZExt(createIntrinsicCall(builder, id, args),
                              builder.getInt32Ty());
  }];
}

def NVVM_BarrierArriveOp : NVVM_VoidIntrinsicOp<"barrier.arrive">
{
  let arguments = (ins Optional<I32>:$barrierId, I32:$numberOfThreads,
                       DefaultValuedAttr<BoolAttr, "true">:$aligned);

  let description = [{
    Thread that executes this op announces their arrival at the barrier with 
    given id and continue their execution.

    The default barrier id is 0 that is similar to `nvvm.barrier` Op. When 
    `barrierId` is not present, the default barrier id is used. 

    The `aligned` attribute, which defaults to `true`, generates the aligned
    form of the barrier (all threads in the CTA execute the same barrier
    instruction). When set to `false`, the unaligned form is generated.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
  }];
  
  let assemblyFormat = "(`id` `=` $barrierId^)? `number_of_threads` `=` $numberOfThreads attr-dict";
}

def NVVM_ClusterArriveOp : NVVM_Op<"cluster.arrive"> {
  let arguments = (ins OptionalAttr<UnitAttr>:$aligned);

  let summary = "Cluster Barrier Arrive Op";
  let description = [{
    The `cluster.arrive` can be used by the threads within the cluster for synchronization and
    communication. The `cluster.arrive` instruction marks the warps' arrival at the barrier
    without causing the executing thread to wait for other participating threads.

    The `aligned` attribute, when provided, generates the .aligned version of the PTX instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-barrier-cluster)
  }];

  string llvmBuilder = [{
      if ($aligned)
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive_aligned);
      else
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive);
  }];
  let assemblyFormat = "attr-dict";
}

def NVVM_ClusterArriveRelaxedOp : NVVM_Op<"cluster.arrive.relaxed", [NVVMRequiresSM<90>]> {
  let arguments = (ins OptionalAttr<UnitAttr>:$aligned);

  let summary = "Cluster Barrier Relaxed Arrive Op";
  let description = [{
    The `cluster.arrive` can be used by the threads within the cluster for synchronization and
    communication. The `cluster.arrive` instruction marks the warps' arrival at the barrier
    without causing the executing thread to wait for other participating threads.

    The `aligned` attribute, when provided, generates the .aligned version of the PTX instruction.
    The .relaxed qualifier on `cluster.arrive` specifies that there are no memory
    ordering and visibility guarantees provided for the memory accesses performed prior to
    `cluster.arrive`.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-barrier-cluster)
  }];

  string llvmBuilder = [{
      if ($aligned)
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive_relaxed_aligned);
      else
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_arrive_relaxed);
  }];
  let assemblyFormat = "attr-dict";
}

def NVVM_ClusterWaitOp : NVVM_Op<"cluster.wait", [NVVMRequiresSM<90>]> {
  let arguments = (ins OptionalAttr<UnitAttr>:$aligned);

  let summary = "Cluster Barrier Wait Op";
  let description = [{
    The `cluster.wait` causes the executing thread to wait for all non-exited threads
    of the cluster to perform `cluster.arrive`. The `aligned` attribute, when provided,
    generates the .aligned version of the PTX instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-barrier-cluster)
  }];

  string llvmBuilder = [{
      if ($aligned)
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_wait_aligned);
      else
        createIntrinsicCall(builder, llvm::Intrinsic::nvvm_barrier_cluster_wait);
  }];
  let assemblyFormat = "attr-dict";
}

//===----------------------------------------------------------------------===//
// NVVM Member/Fence
//===----------------------------------------------------------------------===//

def NVVM_MembarOp : NVVM_Op<"memory.barrier">,
                    Arguments<(ins MemScopeKindAttr:$scope)> {
  let summary = "Memory barrier operation";
  let description = [{
    `membar` operation guarantees that prior memory accesses requested by this
    thread are performed at the specified `scope`, before later memory
    operations requested by this thread following the membar instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "$scope attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(builder, getMembarIntrinsicID($scope));
  }];
}

def NVVM_FenceScClusterOp : NVVM_Op<"fence.sc.cluster"> {
  string llvmBuilder = [{
      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_fence_sc_cluster);
  }];
  let assemblyFormat = "attr-dict";
}

defvar MemOrderAcquireOrRelease =
  ConfinedAttr<MemOrderKindAttr,
    [EnumAttrIsOneOf<MemOrderKindAttr,
      [MemOrderKindAcquire, MemOrderKindRelease]>]>;

def NVVM_FenceSyncRestrictOp : NVVM_Op<"fence.sync_restrict">,
                               Arguments<(ins MemOrderAcquireOrRelease:$order)> {
  let summary = "Uni-directional thread fence operation";
  let description = [{
    The `nvvm.fence.sync_restrict` Op restricts the class of memory
    operations for which the fence instruction provides the memory ordering guarantees.
    `sync_restrict` restricts `acquire` memory semantics to `shared_cluster` and
    `release` memory semantics to `shared_cta` with cluster scope.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(builder, getFenceSyncRestrictID($order));
  }];
}

def NVVM_FenceMbarrierInitOp : NVVM_Op<"fence.mbarrier.init"> {
    let description = [{
    Fence operation that applies on the prior nvvm.mbarrier.init

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_fence_mbarrier_init_release_cluster);
  }];
}

def ProxyAlias : I32EnumAttrCase<"alias", 0, "alias">;
def ProxyAsync   : I32EnumAttrCase<"async", 1, "async">;
def ProxyAsyncGlobal   : I32EnumAttrCase<"async_global", 2, "async.global">;
def ProxyAsyncShared   : I32EnumAttrCase<"async_shared", 3, "async.shared">;
def ProxyTensorMap : I32EnumAttrCase<"TENSORMAP", 4, "tensormap">;
def ProxyGeneric : I32EnumAttrCase<"GENERIC", 5, "generic">;
def ProxyKind : I32EnumAttr<"ProxyKind", "Proxy kind",
  [ProxyAlias, ProxyAsync, ProxyAsyncGlobal, ProxyAsyncShared, ProxyTensorMap, ProxyGeneric]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}

def ProxyKindAttr : EnumAttr<NVVM_Dialect, ProxyKind, "proxy_kind"> {
  let description = [{
    ProxyKind attribute represents a memory proxy which is an abstract label
    applied to a method of memory access. When two memory operations use distinct
    methods of memory access, they are said to be different proxies.
  }];
  let assemblyFormat = "`<` $value `>`";
}

defvar ProxyKindNotTensormapOrGeneric =
  ConfinedAttr<ProxyKindAttr,
    [EnumAttrIsNoneOf<ProxyKindAttr, [ProxyTensorMap, ProxyGeneric]>]>;

def NVVM_FenceProxyOp : NVVM_Op<"fence.proxy">,
  Arguments<(ins ProxyKindNotTensormapOrGeneric:$kind,
                 OptionalAttr<SharedSpaceAttr>:$space)> {
  let description = [{
    Fence operation with proxy to establish an ordering between memory accesses
    that may happen through different proxies.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];
  
  let assemblyFormat = "attr-dict";

  let llvmBuilder = [{
    createIntrinsicCall(builder, getFenceProxyID($kind, $space));
  }];

  let hasVerifier = 1;
}

def NVVM_FenceProxyAcquireOp : NVVM_Op<"fence.proxy.acquire">,
      Arguments<(ins MemScopeKindAttr:$scope, LLVM_PointerGeneric:$addr, I32:$size,
                     DefaultValuedAttr<ProxyKindAttr,
                                       "ProxyKind::GENERIC">:$fromProxy,
                     DefaultValuedAttr<ProxyKindAttr,
                                       "ProxyKind::TENSORMAP">:$toProxy)> {
  let summary = "Uni-directional proxy fence operation with acquire semantics";
  let description = [{
    `fence.proxy.acquire` is a uni-directional fence used to establish ordering
    between a prior memory access performed via the generic proxy and a
    subsequent memory access performed via the tensormap proxy

    The address operand `addr` and the operand `size` together specify the
    memory range `[addr, addr+size)` on which the ordering guarantees on the
    memory accesses across the proxies is to be provided. The only supported
    value for the `size` operand is 128 and must be an immediate. Generic Addressing
    is used unconditionally, and the address specified by the operand `addr` must
    fall within the `.global` state space. Otherwise, the behavior is undefined
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "$scope $addr `,` $size (`from_proxy` `=` $fromProxy^)? (`to_proxy` `=` $toProxy^)? attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(
        builder,
        getUnidirectionalFenceProxyID($fromProxy, $toProxy, $scope, false),
        {$addr, $size});
  }];

  let hasVerifier = 1;
}

def NVVM_FenceProxyReleaseOp : NVVM_Op<"fence.proxy.release">,
      Arguments<(ins MemScopeKindAttr:$scope,
                     DefaultValuedAttr<ProxyKindAttr,
                                       "ProxyKind::GENERIC">:$fromProxy,
                     DefaultValuedAttr<ProxyKindAttr,
                                       "ProxyKind::TENSORMAP">:$toProxy)> {
  let summary = "Uni-directional proxy fence operation with release semantics";
  let description = [{
    `fence.proxy.release` is a uni-directional fence used to establish ordering
    between a prior memory access performed via the generic proxy and a
    subsequent memory access performed via the tensormap proxy. `fence.proxy.release`
    operation can form a release sequence that synchronizes with an acquire
    sequence that contains the fence.proxy.acquire proxy fence operation
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "$scope (`from_proxy` `=` $fromProxy^)? (`to_proxy` `=` $toProxy^)? attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(builder, getUnidirectionalFenceProxyID(
                                     $fromProxy, $toProxy, $scope, true));
  }];

  let hasVerifier = 1;
}

def NVVM_FenceProxySyncRestrictOp : NVVM_Op<"fence.proxy.sync_restrict">,
      Arguments<(ins MemOrderAcquireOrRelease:$order,
                     DefaultValuedAttr<ProxyKindAttr, "ProxyKind::GENERIC">:$fromProxy,
                     DefaultValuedAttr<ProxyKindAttr, "ProxyKind::async">:$toProxy)> {
  let summary = "Uni-directional proxy fence operation with sync_restrict";
  let description = [{
    The `nvvm.fence.proxy.sync_restrict` Op used to establish
    ordering between a prior memory access performed between proxies. Currently,
    the ordering is only supported between async and generic proxies. `sync_restrict`
    restricts `acquire` memory semantics to `shared_cluster` and `release` memory
    semantics to `shared_cta` with cluster scope.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-membar)
  }];

  let assemblyFormat = "attr-dict";
  let llvmBuilder = [{
    createIntrinsicCall(builder, getFenceProxySyncRestrictID($order));
  }];

  let hasVerifier = 1;
}

def SetMaxRegisterActionIncrease : I32EnumAttrCase<"increase", 0>;
def SetMaxRegisterActionDecrease   : I32EnumAttrCase<"decrease", 1>;
def SetMaxRegisterAction : I32EnumAttr<"SetMaxRegisterAction", "NVVM set max register action",
  [SetMaxRegisterActionDecrease, SetMaxRegisterActionIncrease]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def SetMaxRegisterActionAttr : EnumAttr<NVVM_Dialect, SetMaxRegisterAction, "action">;

def NVVM_SetMaxRegisterOp : NVVM_Op<"setmaxregister"> {
  let arguments = (ins I32Attr:$regCount, SetMaxRegisterActionAttr:$action);
  let assemblyFormat = "$action $regCount attr-dict";
  let hasVerifier = 1;
  string llvmBuilder = [{
    auto intId = (op.getAction() == NVVM::SetMaxRegisterAction::increase) ?
      llvm::Intrinsic::nvvm_setmaxnreg_inc_sync_aligned_u32 :
      llvm::Intrinsic::nvvm_setmaxnreg_dec_sync_aligned_u32;

    createIntrinsicCall(builder, intId, builder.getInt32($regCount));
  }];
}

def ShflKindBfly : I32EnumAttrCase<"bfly", 0>;
def ShflKindUp   : I32EnumAttrCase<"up", 1>;
def ShflKindDown : I32EnumAttrCase<"down", 2>;
def ShflKindIdx  : I32EnumAttrCase<"idx", 3>;

/// Enum attribute of the different shuffle kinds.
def ShflKind : I32EnumAttr<"ShflKind", "NVVM shuffle kind",
  [ShflKindBfly, ShflKindUp, ShflKindDown, ShflKindIdx]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def ShflKindAttr : EnumAttr<NVVM_Dialect, ShflKind, "shfl_kind">;

def NVVM_ShflOp :
  NVVM_Op<"shfl.sync", [NVVMRequiresSM<30>, InferTypeOpAdaptor]>,
  Results<(outs AnyTypeOf<[I32, F32, LLVMStructType]>:$res)>,
  Arguments<(ins I32:$thread_mask,
                 AnyTypeOf<[I32, F32]>:$val,
                 I32:$offset,
                 I32:$mask_and_clamp,
                 ShflKindAttr:$kind,
                 OptionalAttr<UnitAttr>:$return_value_and_is_valid)> {
  let summary = "NVVM Dialect Op for shfl.sync";
  let description = [{
    The `shfl.sync` Op implements data shuffle within threads of a warp.
    The `thread_mask` denotes the threads participating in the Op where
    the bit position corresponds to a particular thread's laneid.
    The `offset` specifies a source lane or source lane offset
    (depending on `kind`). The `val` is the input value to be copied from
    the source. The `mask_and_clamp` contains two packed values specifying
    a mask for logically splitting warps into sub-segments and an upper bound
    for clamping the source lane index.
    
    The `return_value_and_is_valid` unit attribute can be specified to indicate 
    that the return value is a two-element struct, where the first element is 
    the result value and the second element is a predicate indicating if the 
    computed source lane index is valid.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync)
  }];
  string llvmBuilder = [{
      auto intId = getShflIntrinsicId(
          $_resultType, $kind, static_cast<bool>($return_value_and_is_valid));
      $res = createIntrinsicCall(builder,
          intId, {$thread_mask, $val, $offset, $mask_and_clamp});
  }];
  let assemblyFormat = [{
    $kind $thread_mask `,` $val `,` $offset `,` $mask_and_clamp  attr-dict
     `:` type($val) `->` type($res)
   }];
   let hasVerifier = 1;
}

def VoteSyncKindAny : I32EnumAttrCase<"any", 0>;
def VoteSyncKindAll : I32EnumAttrCase<"all", 1>;
def VoteSyncKindBallot : I32EnumAttrCase<"ballot", 2>;
def VoteSyncKindUni : I32EnumAttrCase<"uni", 3>;

def VoteSyncKind : I32EnumAttr<"VoteSyncKind", "NVVM vote sync kind",
                               [VoteSyncKindAny, VoteSyncKindAll,
                                VoteSyncKindBallot, VoteSyncKindUni]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}

def VoteSyncKindAttr : EnumAttr<NVVM_Dialect, VoteSyncKind, "vote_sync_kind">;

def NVVM_VoteSyncOp
    : NVVM_Op<"vote.sync", [InferTypeOpAdaptor]>,
      Results<(outs AnyTypeOf<[I32, I1]>:$res)>,
      Arguments<(ins I32:$mask, I1:$pred, VoteSyncKindAttr:$kind)> {
  let summary = "Vote across thread group";
  let description = [{
    The `vote.sync` op will cause executing thread to wait until all non-exited
    threads corresponding to membermask have executed `vote.sync` with the same
    qualifiers and same membermask value before resuming execution.

    The vote operation kinds are:
    - `any`: True if source predicate is True for some thread in membermask.
    - `all`: True if source predicate is True for all non-exited threads in
      membermask. 
    - `uni`: True if source predicate has the same value in all non-exited
      threads in membermask.
    - `ballot`: In the ballot form, the destination result is a 32 bit integer.
      In this form, the predicate from each thread in membermask are copied into
      the corresponding bit position of the result, where the bit position
      corresponds to the thread's lane id.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync)
  }];
  string llvmBuilder = [{
    auto intId = getVoteSyncIntrinsicId($kind);
    $res = createIntrinsicCall(builder, intId, {$mask, $pred});
  }];
  let assemblyFormat = "$kind $mask `,` $pred attr-dict `->` type($res)";
  let hasVerifier = 1;
}

def NVVM_SyncWarpOp :
  NVVM_IntrOp<"bar.warp.sync">,
  Arguments<(ins I32:$mask)> {
  let summary = "Warp Barrier Synchronization Op";
  let description = [{
    The `nvvm.bar.warp.sync` operation performs barrier synchronization for threads 
    within a warp. 

    This operation causes the executing thread to wait until all threads corresponding 
    to the `mask` operand have executed a `bar.warp.sync` with the same mask value 
    before resuming execution.

    The `mask` operand specifies the threads participating in the barrier, where each 
    bit position corresponds to the thread's lane ID within the warp. Only threads with 
    their corresponding bit set in the mask participate in the barrier synchronization.

    **Important constraints**:
    - The behavior is undefined if the executing thread is not included in the mask 
      (i.e., the bit corresponding to the thread's lane ID is not set)
    - For compute capability sm_6x or below, all threads in the mask must execute 
      the same `bar.warp.sync` instruction in convergence

    This operation also guarantees memory ordering among participating threads. 
    Threads within the warp that wish to communicate via memory can store to memory, 
    execute `bar.warp.sync`, and then safely read values stored by other threads 
    in the warp.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar-warp-sync)
  }];

  let assemblyFormat = "$mask attr-dict `:` type($mask)";
}

def NVVM_ElectSyncOp : NVVM_Op<"elect.sync">
{  
  let summary = "Elect one leader thread";
  let description = [{
    The `elect.sync` instruction elects one predicated active leader
    thread from among a set of threads specified in the `membermask`.
    When the `membermask` is not provided explicitly, a default value
    of `0xFFFFFFFF` is used. The predicate result is set to `True` for
    the leader thread, and `False` for all other threads.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-elect-sync)
  }];

  let arguments = (ins Optional<I32>:$membermask);
  let results = (outs I1:$pred);
  let assemblyFormat = "($membermask^)? attr-dict `->` type(results)";
  string llvmBuilder = [{
    auto *resultTuple = createIntrinsicCall(builder,
        llvm::Intrinsic::nvvm_elect_sync,
        {$membermask ? $membermask : builder.getInt32(0xFFFFFFFF)});
    // Extract the second value into $pred
    $pred = builder.CreateExtractValue(resultTuple, 1);
  }];
}

//===----------------------------------------------------------------------===//
// Permute Bytes (Prmt)
//===----------------------------------------------------------------------===//

// Attributes for the permute operation modes supported by PTX.
def PermuteModeDefault : I32EnumAttrCase<"DEFAULT", 0, "default">;
def PermuteModeF4E : I32EnumAttrCase<"F4E", 1, "f4e">;
def PermuteModeB4E : I32EnumAttrCase<"B4E", 2, "b4e">;
def PermuteModeRC8 : I32EnumAttrCase<"RC8", 3, "rc8">;
def PermuteModeECL : I32EnumAttrCase<"ECL", 4, "ecl">;
def PermuteModeECR : I32EnumAttrCase<"ECR", 5, "ecr">;
def PermuteModeRC16 : I32EnumAttrCase<"RC16", 6, "rc16">;

def PermuteMode : I32EnumAttr<"PermuteMode", "NVVM permute mode",
                              [PermuteModeDefault, PermuteModeF4E,
                               PermuteModeB4E, PermuteModeRC8, PermuteModeECL,
                               PermuteModeECR, PermuteModeRC16]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}

def PermuteModeAttr : EnumAttr<NVVM_Dialect, PermuteMode, "permute_mode"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_PermuteOp : NVVM_SingleResultIntrinsicOp<"prmt", [Pure]>,
                     Results<(outs I32:$res)>,
                     Arguments<(ins I32:$lo, Optional<I32>:$hi, I32:$selector,
                         PermuteModeAttr:$mode)> {
  let summary = "Permute bytes from two 32-bit registers";
  let description = [{
    The `nvvm.prmt` operation constructs a permutation of the
    bytes of the first one or two operands, selecting based on 
    the 2 least significant bits of the final operand.
    
    The bytes in the first one or two source operands are numbered. 
    The first source operand (%lo) is numbered {b3, b2, b1, b0}, 
    in the case of the '``default``', '``f4e``' and '``b4e``' variants, 
    the second source operand (%hi) is numbered {b7, b6, b5, b4}.

    Modes:
    - `default`: Index mode         - each nibble in `selector` selects a byte from the 8-byte pool
    - `f4e`    : Forward 4 extract  - extracts 4 contiguous bytes starting from position in `selector`
    - `b4e`    : Backward 4 extract - extracts 4 contiguous bytes in reverse order
    - `rc8`    : Replicate 8        - replicates the lower 8 bits across the 32-bit result
    - `ecl`    : Edge clamp left    - clamps out-of-range indices to the leftmost valid byte
    - `ecr`    : Edge clamp right   - clamps out-of-range indices to the rightmost valid byte  
    - `rc16`   : Replicate 16       - replicates the lower 16 bits across the 32-bit result

    Depending on the 2 least significant bits of the %selector operand, the result
    of the permutation is defined as follows:

    +------------+----------------+--------------+
    |    Mode    | %selector[1:0] |    Output    |
    +------------+----------------+--------------+
    | '``f4e``'  | 0              | {3, 2, 1, 0} |
    |            +----------------+--------------+
    |            | 1              | {4, 3, 2, 1} |
    |            +----------------+--------------+
    |            | 2              | {5, 4, 3, 2} |
    |            +----------------+--------------+
    |            | 3              | {6, 5, 4, 3} |
    +------------+----------------+--------------+
    | '``b4e``'  | 0              | {5, 6, 7, 0} |
    |            +----------------+--------------+
    |            | 1              | {6, 7, 0, 1} |
    |            +----------------+--------------+
    |            | 2              | {7, 0, 1, 2} |
    |            +----------------+--------------+
    |            | 3              | {0, 1, 2, 3} |
    +------------+----------------+--------------+
    | '``rc8``'  | 0              | {0, 0, 0, 0} |
    |            +----------------+--------------+
    |            | 1              | {1, 1, 1, 1} |
    |            +----------------+--------------+
    |            | 2              | {2, 2, 2, 2} |
    |            +----------------+--------------+
    |            | 3              | {3, 3, 3, 3} |
    +------------+----------------+--------------+
    | '``ecl``'  | 0              | {3, 2, 1, 0} |
    |            +----------------+--------------+
    |            | 1              | {3, 2, 1, 1} |
    |            +----------------+--------------+
    |            | 2              | {3, 2, 2, 2} |
    |            +----------------+--------------+
    |            | 3              | {3, 3, 3, 3} |
    +------------+----------------+--------------+
    | '``ecr``'  | 0              | {0, 0, 0, 0} |
    |            +----------------+--------------+
    |            | 1              | {1, 1, 1, 0} |
    |            +----------------+--------------+
    |            | 2              | {2, 2, 1, 0} |
    |            +----------------+--------------+
    |            | 3              | {3, 2, 1, 0} |
    +------------+----------------+--------------+
    | '``rc16``' | 0              | {1, 0, 1, 0} |
    |            +----------------+--------------+
    |            | 1              | {3, 2, 3, 2} |
    |            +----------------+--------------+
    |            | 2              | {1, 0, 1, 0} |
    |            +----------------+--------------+
    |            | 3              | {3, 2, 3, 2} |
    +------------+----------------+--------------+
    
    [For more information, see PTX ISA]
    (https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prmt)
  }];

  let assemblyFormat = [{
    $mode $selector `,` $lo  (`,` $hi^)?  attr-dict `:` type($res)
  }];

  let hasVerifier = 1;
}

def LoadCacheModifierCA : I32EnumAttrCase<"CA", 0, "ca">;
def LoadCacheModifierCG : I32EnumAttrCase<"CG", 1, "cg">;
def LoadCacheModifierCS : I32EnumAttrCase<"CS", 2, "cs">;
def LoadCacheModifierLU : I32EnumAttrCase<"LU", 3, "lu">;
def LoadCacheModifierCV : I32EnumAttrCase<"CV", 4, "cv">;

/// Enum attribute of the different kinds.
def LoadCacheModifierKind : I32EnumAttr<"LoadCacheModifierKind", 
                                "NVVM load cache modifier kind",
  [LoadCacheModifierCA, LoadCacheModifierCG, LoadCacheModifierCS, 
    LoadCacheModifierLU, LoadCacheModifierCV]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
  let description = [{
    Enum attribute of the different kinds of cache operators for load instructions.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#id62)    
  }];
}

def LoadCacheModifierAttr : EnumAttr<NVVM_Dialect, LoadCacheModifierKind, "load_cache_modifier">;

def NVVM_CpAsyncOp : NVVM_Op<"cp.async.shared.global">,
  Arguments<(ins LLVM_PointerShared:$dst,
                 LLVM_PointerGlobal:$src,
                 I32Attr:$size,
                 LoadCacheModifierAttr:$modifier,
                 Optional<I32>:$cpSize)> {
  let assemblyFormat = "$dst `,` $src `,` $size `,` `cache` `=` $modifier (`,` $cpSize^)? attr-dict `:` type(operands)";
  let hasVerifier = 1;
  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::SmallVector<llvm::Value *> &args);
  }];
  string llvmBuilder = [{
    llvm::SmallVector<llvm::Value *> translatedOperands;
    auto id = NVVM::CpAsyncOp::getIntrinsicIDAndArgs(
      *op, moduleTranslation, translatedOperands);
    createIntrinsicCall(builder, id, translatedOperands);
  }];
}

def NVVM_CpAsyncCommitGroupOp : NVVM_Op<"cp.async.commit.group"> {
  string llvmBuilder = [{
      createIntrinsicCall(builder, llvm::Intrinsic::nvvm_cp_async_commit_group);
  }];
  let assemblyFormat = "attr-dict";
}

def NVVM_CpAsyncWaitGroupOp : NVVM_Op<"cp.async.wait.group">,
  Arguments<(ins I32Attr:$n)> {
  string llvmBuilder = [{
      createIntrinsicCall(
        builder,
        llvm::Intrinsic::nvvm_cp_async_wait_group,
        llvm::ConstantInt::get(
          llvm::Type::getInt32Ty(moduleTranslation.getLLVMContext()),
          $n));
  }];
  let assemblyFormat = "$n attr-dict";
}

def NVVM_CpAsyncMBarrierArriveOp : NVVM_VoidIntrinsicOp<"cp.async.mbarrier.arrive"> {
  let summary = "NVVM Dialect Op for cp.async.mbarrier.arrive";
  let description = [{
    The `cp.async.mbarrier.arrive` Op makes the *mbarrier object* track
    all prior cp.async operations initiated by the executing thread.
    The `addr` operand specifies the address of the *mbarrier object*
    in generic or shared::cta address space. When it is generic, the
    underlying memory should fall within the shared::cta space;
    otherwise the behavior is undefined. The `noinc` attr impacts
    how the mbarrier's state is updated.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr,
    DefaultValuedAttr<I1Attr, "0">:$noinc);

  let assemblyFormat = "$addr attr-dict `:` type(operands)";
}

//===----------------------------------------------------------------------===//
// NVVM Conversion Ops (for "cvt.*" family of PTX instructions)
//===----------------------------------------------------------------------===//

def NVVM_ConvertFloatToTF32Op : NVVM_Op<"convert.float.to.tf32"> {
  let summary = "Convert the given float input to TF32";
  let description = [{
    This Op converts the given f32 input to tf32.
    The result `res` is represented as an i32 type.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction. The `rnd` and `sat` attributes specify the
    the rounding and saturation modes respectively.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let hasVerifier = 1;
  let results = (outs I32:$res);
  let arguments = (ins
    F32:$src,
    DefaultValuedAttr<FPRoundingModeAttr, "FPRoundingMode::NONE">:$rnd,
    DefaultValuedAttr<SaturationModeAttr, "SaturationMode::NONE">:$sat,
    DefaultValuedAttr<BoolAttr, "false">:$relu);

  let assemblyFormat = "$src attr-dict";

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(NVVM::FPRoundingMode,
                                              NVVM::SaturationMode,
                                              bool hasRelu);
  }];

  string llvmBuilder = [{
    auto intId = NVVM::ConvertFloatToTF32Op::getIntrinsicID($rnd, $sat, $relu);
    $res = createIntrinsicCall(builder, intId, {$src});
  }];
}

def NVVM_ConvertF32x2ToF4x2Op : NVVM_Op<"convert.f32x2.to.f4x2"> {
  let summary = "Convert a pair of float inputs to f4x2";
  let description = [{
    This Op converts each of the given float inputs to the specified fp4 type.
    The result `dst` is returned as an i8 type where the converted values are 
    packed such that the value converted from `a` is stored in the upper 4 bits 
    of `dst` and the value converted from `b` is stored in the lower 4 bits of 
    `dst`.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let results = (outs I8:$dst);
  let arguments = (ins F32:$a, F32:$b,
                       DefaultValuedAttr<BoolAttr, "false">:$relu,
                       TypeAttr:$dstTy);
  let assemblyFormat = "$a `,` $b attr-dict `:` type($dst) `(` $dstTy `)`";
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(NVVM::ConvertF32x2ToF4x2Op op, 
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];

  string llvmBuilder = [{
    auto [intId, args] = NVVM::ConvertF32x2ToF4x2Op::getIntrinsicIDAndArgs(op, moduleTranslation, builder);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, args);
    $dst = builder.CreateTruncOrBitCast(packedI16, llvm::Type::getInt8Ty(builder.getContext()));
  }];
}

class NVVM_ConvertFPx2ToF4x2Op<string srcType>
    : NVVM_Op<"convert."#!tolower(srcType)#"x2.to.f4x2"> {
  let summary = "Convert an " # !tolower(srcType) # "x2 input to f4x2";
  let description = [{
    This Op converts each of the given }]#srcType#[{ inputs in an }]#!tolower
    (srcType)#[{x2 vector to the specified fp4 type.
    The result `dst` is returned as an i8 type where the converted values are 
    packed such that the value converted from the first element of `a` is 
    stored in the lower 4 bits of `dst` and the value converted from the second 
    element of `a` is stored in the upper 4 bits of `dst`.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.
    }];
  let results = (outs I8:$dst);
  let arguments = (ins
      VectorOfLengthAndType<[2], [!cast<Type>(srcType)]>:$src,
      DefaultValuedAttr<BoolAttr, "false">:$relu,
      TypeAttrOf<F4E2M1FN>:$dstTy);

  let assemblyFormat =
      "$src attr-dict `:` type($src) `->` type($dst) `(` $dstTy `)`";

  let extraClassDeclaration = [{
    static NVVM::IDArgPair
    getIntrinsicIDAndArgs(NVVM::Convert}]#srcType#[{x2ToF4x2Op &op, 
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];

  string llvmBuilder = [{
    auto [intId, args] = NVVM::Convert}]#srcType#[{x2ToF4x2Op::getIntrinsicIDAndArgs(op, moduleTranslation, builder);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, args);
    $dst = builder.CreateTruncOrBitCast(packedI16, llvm::Type::getInt8Ty(builder.getContext()));
  }];
}

def NVVM_ConvertF16x2ToF4x2Op : NVVM_ConvertFPx2ToF4x2Op<"F16">;
def NVVM_ConvertBF16x2ToF4x2Op : NVVM_ConvertFPx2ToF4x2Op<"BF16">;

def NVVM_ConvertF32x2ToF6x2Op : NVVM_Op<"convert.f32x2.to.f6x2"> {
  let summary = "Convert a pair of float inputs to f6x2";
  let description = [{
    This Op converts each of the given float inputs to the specified fp6 type.
    The result `dst` is represented either as an i16 type or as a vector
    of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from `a` is stored in the upper 8 bits of `dst` 
    with 2 MSB bits padded with zeros and the value converted from `b` is 
    stored in the lower 8 bits of `dst` with 2 MSB bits padded with zeros.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins 
    F32:$a,
    F32:$b,
    DefaultValuedAttr<BoolAttr, "false">:$relu,
    TypeAttr:$dstTy);
  let assemblyFormat = "$a `,` $b attr-dict `:` type($dst) `(` $dstTy `)`";
  let hasVerifier = 1;
  
  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(mlir::Type dstTy,
                                              bool hasRelu);
  }];

  string llvmBuilder = [{
    auto intId = NVVM::ConvertF32x2ToF6x2Op::getIntrinsicID($dstTy, $relu);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, {$a, $b});
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

class NVVM_ConvertFPx2ToF6x2Op<string srcType>
    : NVVM_Op<"convert."#!tolower(srcType)#"x2.to.f6x2"> {
  let summary = "Convert an " # !tolower(srcType) # "x2 input to f6x2";
  let description = [{
    This Op converts each of the given }]#srcType#[{ inputs in an }]#!tolower
    (srcType)#[{x2 vector to the specified fp6 type. The result `dst` is 
    represented either as an i16 type or as a vector of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from the first element of `a` is stored in the 
    lower 8 bits of `dst` with 2 MSB bits padded with zeros and the value 
    converted from the second element of `a` is stored in the upper 8 bits of 
    `dst` with 2 MSB bits padded with zeros.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector with 2 MSB bits padded with zeros.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.
  }];

  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins
      VectorOfLengthAndType<[2], [!cast<Type>(srcType)]>:$src,
      DefaultValuedAttr<BoolAttr, "false">:$relu,
      TypeAttrOf<AnyTypeOf<[F6E2M3FN, F6E3M2FN]>>:$dstTy);
  let assemblyFormat = "$src attr-dict `:` type($src) `->` type($dst) `(` $dstTy `)`";
  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(mlir::Type dstTy, bool hasRelu);
  }];

  string llvmBuilder = [{
    auto intId = NVVM::Convert}]#srcType#[{x2ToF6x2Op::getIntrinsicID($dstTy, $relu);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, {$src});
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                        llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def NVVM_ConvertF16x2ToF6x2Op : NVVM_ConvertFPx2ToF6x2Op<"F16">;
def NVVM_ConvertBF16x2ToF6x2Op : NVVM_ConvertFPx2ToF6x2Op<"BF16">;

def NVVM_ConvertF32x2ToF8x2Op : NVVM_Op<"convert.f32x2.to.f8x2"> {
  let summary = "Convert a pair of float inputs to f8x2";
  let description = [{
    This Op converts each of the given float inputs to the specified fp8 type.
    The result `dst` is represented as an i16 type or as a vector
    of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from `a` is stored in the upper 8 bits of `dst` 
    and the value converted from `b` is stored in the lower 8 bits of `dst`.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `rnd` and `sat` attributes specify the rounding and saturation modes respectively.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let hasVerifier = 1;
  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins
    F32:$a,
    F32:$b,
    DefaultValuedAttr<FPRoundingModeAttr, "FPRoundingMode::NONE">:$rnd,
    DefaultValuedAttr<SaturationModeAttr, "SaturationMode::NONE">:$sat,
    DefaultValuedAttr<BoolAttr, "false">:$relu,
    TypeAttr:$dstTy);
  let assemblyFormat = "$a `,` $b attr-dict `:` type($dst) `(` $dstTy `)`";

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(mlir::Type dstTy,
                                              NVVM::FPRoundingMode rnd,
                                              NVVM::SaturationMode sat,
                                              bool hasRelu);
  }];
  
  string llvmBuilder = [{
    auto intId = NVVM::ConvertF32x2ToF8x2Op::getIntrinsicID($dstTy, $rnd, $sat, $relu);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, {$a, $b});
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def NVVM_ConvertF16x2ToF8x2Op : NVVM_Op<"convert.f16x2.to.f8x2"> {
  let summary = "Convert an f16x2 input to f8x2";
  let description = [{
    This Op converts the given f16 inputs in an f16x2 vector to the specified 
    f8 type.
    The result `dst` is represented as an i16 type or as a vector
    of two i8 types.
    If `dst` is returned as an i16 type, the converted values from `a`
    are packed such that the value converted from the first element of `a`
    is stored in the upper 8 bits of `dst` and the value converted from the
    second element of `a` is stored in the lower 8 bits of `dst`.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let hasVerifier = 1;
  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins
    VectorOfLengthAndType<[2], [F16]>:$a,
    DefaultValuedAttr<BoolAttr, "false">:$relu,
    TypeAttr:$dstTy);
  let assemblyFormat = "$a attr-dict `:` type($a) `->` type($dst) `(` $dstTy `)`";

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(mlir::Type dstTy,
                                              bool hasRelu);
  }];

  string llvmBuilder = [{
    auto intId = NVVM::ConvertF16x2ToF8x2Op::getIntrinsicID($dstTy, $relu);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, {$a});
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def NVVM_ConvertBF16x2ToF8x2Op : NVVM_Op<"convert.bf16x2.to.f8x2"> {
  let summary = "Convert a pair of bf16 inputs to f8x2";
  let description = [{
    This Op converts the given bf16 inputs in a bf16x2 vector to the specified 
    f8 type. The result `dst` is represented either as a packed i16 type or as 
    a vector of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from the first element of `a` is stored in the 
    lower 8 bits of `dst` and the value converted from the second element of 
    `a` is stored in the upper 8 bits of `dst`.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `rnd` and `sat` attributes specify the rounding and saturation modes 
    respectively.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let hasVerifier = 1;
  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins
    VectorOfLengthAndType<[2], [BF16]>:$src,
    DefaultValuedAttr<FPRoundingModeAttr, "FPRoundingMode::NONE">:$rnd,
    DefaultValuedAttr<SaturationModeAttr, "SaturationMode::NONE">:$sat,
    DefaultValuedAttr<BoolAttr, "false">:$relu,
    TypeAttrOf<AnyTypeOf<[F8E8M0FNU, F8E4M3FN, F8E5M2]>>:$dstTy);
  let assemblyFormat = "$src attr-dict `:` type($src) `->` type($dst) `(` $dstTy `)`";
  
  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(mlir::Type dstTy,
                                  NVVM::FPRoundingMode rnd,
                                  NVVM::SaturationMode sat, bool hasRelu);
  }];

  string llvmBuilder = [{
    auto intId = NVVM::ConvertBF16x2ToF8x2Op::getIntrinsicID($dstTy, $rnd, $sat, $relu);
    llvm::Value *packedI16 = createIntrinsicCall(builder, intId, {$src});
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def SaturationModeSatfiniteOrNone :
  ConfinedAttr<SaturationModeAttr, [EnumAttrIsOneOf<SaturationModeAttr,
                [SaturationModeNone, SaturationModeFinite]>]>;

class NVVM_ConvertToFP16x2Op_Base <string srcTypeStr, Type srcStorageType, string dstTypeStr, list<Type> supportedTypes, int needVerify = 0>
: NVVM_SingleResultIntrinsicOp<"convert." # !tolower(srcTypeStr) # "x2.to." # !tolower(dstTypeStr) # "x2", [], "$dst"> {
  let summary = "Convert a pair of " # !tolower(srcTypeStr) # " inputs to " # !tolower(dstTypeStr) # "x2";
  let description = [{
    This Op converts the given }] # !tolower(srcTypeStr) # [{ inputs in a }] #
    !if(!eq(srcTypeStr, "F4"), "packed i8", "i8x2 vector") # [{ to }] #
    !tolower(dstTypeStr) # [{.

    The result `dst` is represented as a vector of }] # !tolower(dstTypeStr) # [{ elements.

    The `relu` attribute, when set, lowers to the '.relu' variant of
    the cvt instruction.}] #

    !if(!eq(dstTypeStr, "BF16"),
    [{

    The `sat` attribute specifies the saturation mode.

    The optional scaling-factors for each of the inputs are provided through
    the operand `scaleFactor` as a packed i16 type. Only `ue8m0` is supported
    as the type of the scale-factor currently.}], "") # [{

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];
  let results = (outs VectorOfLengthAndType<[2], [!cast<Type>(dstTypeStr)]>:$dst);
  let arguments = !if(!eq(dstTypeStr, "F16"),
    (ins srcStorageType:$src,
         TypeAttrOf<AnyTypeOf<supportedTypes>>:$srcType,
         DefaultValuedAttr<BoolAttr, "false">:$relu),
    (ins srcStorageType:$src,
         Optional<I16>:$scaleFactor,
         TypeAttrOf<AnyTypeOf<supportedTypes>>:$srcType,
         DefaultValuedAttr<SaturationModeSatfiniteOrNone, "SaturationMode::NONE">:$sat,
         DefaultValuedAttr<BoolAttr, "false">:$relu));
  let assemblyFormat = 
    !if(!eq(dstTypeStr, "F16"),
      "$src attr-dict `:` type($src) `(` $srcType `)` `->` type($dst)",
      "$src (`,` $scaleFactor^)? attr-dict `:` type($src) `(` $srcType `)` `->` type($dst)");
  let hasVerifier = needVerify;
}

def NVVM_ConvertF8x2ToF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F8", VectorOfLengthAndType<[2], [I8]>, "F16", 
    [F8E4M3FN, F8E5M2]>;
def NVVM_ConvertF8x2ToBF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F8", VectorOfLengthAndType<[2], [I8]>, "BF16",
    [F8E8M0FNU, F8E4M3FN, F8E5M2], 1> {
  let append description = [{

    Example:

    ```mlir
    // Basic conversion from f8E4M3FN.
    %res1 = nvvm.convert.f8x2.to.bf16x2 %src
        : vector<2xi8> (f8E4M3FN) -> vector<2xbf16>

    // Conversion from f8E5M2 with relu and saturation.
    %res2 = nvvm.convert.f8x2.to.bf16x2 %src
        {relu = true, sat = #nvvm.sat_mode<satfinite>}
        : vector<2xi8> (f8E5M2) -> vector<2xbf16>

    // Conversion with a packed ue8m0 scale-factor.
    %res3 = nvvm.convert.f8x2.to.bf16x2 %src, %scaleFactor
        : vector<2xi8> (f8E4M3FN) -> vector<2xbf16>
    ```
  }];
}
def NVVM_ConvertF6x2ToF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F6", VectorOfLengthAndType<[2], [I8]>, "F16", 
    [F6E2M3FN, F6E3M2FN]>;
def NVVM_ConvertF6x2ToBF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F6", VectorOfLengthAndType<[2], [I8]>, "BF16",
    [F6E2M3FN, F6E3M2FN]> {
  let append description = [{

    Example:

    ```mlir
    // Basic conversion from f6E2M3FN.
    %res1 = nvvm.convert.f6x2.to.bf16x2 %src
        : vector<2xi8> (f6E2M3FN) -> vector<2xbf16>

    // Conversion from f6E3M2FN with relu and saturation.
    %res2 = nvvm.convert.f6x2.to.bf16x2 %src
        {relu = true, sat = #nvvm.sat_mode<satfinite>}
        : vector<2xi8> (f6E3M2FN) -> vector<2xbf16>

    // Conversion with a packed ue8m0 scale-factor.
    %res3 = nvvm.convert.f6x2.to.bf16x2 %src, %scaleFactor
        : vector<2xi8> (f6E2M3FN) -> vector<2xbf16>
    ```
  }];
}
def NVVM_ConvertF4x2ToF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F4", I8, "F16", [F4E2M1FN]>;
def NVVM_ConvertF4x2ToBF16x2Op :
  NVVM_ConvertToFP16x2Op_Base<"F4", I8, "BF16", [F4E2M1FN]> {
  let append description = [{

    Example:

    ```mlir
    // Basic conversion; the f4x2 source is packed in a single i8.
    %res1 = nvvm.convert.f4x2.to.bf16x2 %src
        : i8 (f4E2M1FN) -> vector<2xbf16>

    // Conversion with relu and saturation.
    %res2 = nvvm.convert.f4x2.to.bf16x2 %src
        {relu = true, sat = #nvvm.sat_mode<satfinite>}
        : i8 (f4E2M1FN) -> vector<2xbf16>

    // Conversion with a packed ue8m0 scale-factor.
    %res3 = nvvm.convert.f4x2.to.bf16x2 %src, %scaleFactor
        : i8 (f4E2M1FN) -> vector<2xbf16>
    ```
  }];
}

def NVVM_ConvertF32x2ToS2F6x2Op : NVVM_Op<"convert.f32x2.to.s2f6x2"> {
  let summary = "Convert a pair of f32 inputs to S2F6x2";
  let description = [{
    This Op converts each of the given f32 inputs to the
    S2F6x2 type. The result `dst` can be either a packed i16 type or a vector
    of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from `a` is stored in the upper 8 bits of `dst` 
    and the value converted from `b` is stored in the lower 8 bits of `dst`.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `relu` attribute, when set, lowers to the '.relu' variant
    of the cvt instruction.
    The optional scaling-factors for each of the inputs are provided through 
    the operand `scaleFactor` as a packed i16 type. Only `ue8m0` is supported 
    as the type of the scale-factor currently.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins F32:$a, F32:$b,
      Optional<I16>:$scaleFactor,
      DefaultValuedAttr<BoolAttr, "false">:$relu);
  let assemblyFormat =
      "$a `,` $b (`,` $scaleFactor^)? attr-dict `:` type($dst)";
  let extraClassDeclaration = [{
    static IDArgPair getIntrinsicIDAndArgs(Operation &op,
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::ConvertF32x2ToS2F6x2Op::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
    llvm::Value *packedI16 = createIntrinsicCall(builder, id, args);
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def NVVM_ConvertBF16x2ToS2F6x2Op : NVVM_Op<"convert.bf16x2.to.s2f6x2"> {
  let summary = "Convert a pair of BF16 inputs to S2F6x2";
  let description = [{
    This Op converts each of the given BF16 inputs in a bf16x2 vector to the
    S2F6x2 type. The result `dst` can be either a packed i16 type or a vector
    of two i8 types.
    If `dst` is returned as an i16 type, the converted values are packed such 
    that the value converted from the first element of `a` is stored in the 
    lower 8 bits of `dst` and the value converted from the second element of 
    `a` is stored in the upper 8 bits of `dst`.
    If `dst` is returned as a vector type, each converted value is stored as an 
    i8 element in the vector.
    The `relu` attribute, when set, lowers to the '.relu' variant
    of the cvt instruction.
    The optional scaling-factors for each of the inputs are provided through 
    the operand `scaleFactor` as a packed i16 type. Only `ue8m0` is supported 
    as the type of the scale-factor currently.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let results = (outs AnyTypeOf<[I16, VectorOfLengthAndType<[2], [I8]>]>:$dst);
  let arguments = (ins
      VectorOfLengthAndType<[2], [BF16]>:$src,
      Optional<I16>:$scaleFactor,
      DefaultValuedAttr<BoolAttr, "false">:$relu);
  let assemblyFormat =
      "$src (`,` $scaleFactor^)? attr-dict `:` type($src) `->` type($dst)";
  let extraClassDeclaration = [{
    static IDArgPair getIntrinsicIDAndArgs(Operation &op,
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::ConvertBF16x2ToS2F6x2Op::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
    llvm::Value *packedI16 = createIntrinsicCall(builder, id, args);
    if(op.getDst().getType().isInteger(16))
      $dst = packedI16;
    else
      $dst = builder.CreateBitCast(packedI16,
                      llvm::FixedVectorType::get(llvm::Type::getInt8Ty(builder.getContext()), 2));
  }];
}

def NVVM_ConvertS2F6x2ToBF16x2Op : NVVM_SingleResultIntrinsicOp<"convert.s2f6x2.to.bf16x2", [], "$dst"> {
  let summary = "Convert s2f6x2 to bf16x2";
  let description = [{
    This Op converts a pair of s2f6x2 inputs to bf16x2 type. The result `dst` 
    is represented as a vector of two bf16 elements.

    The `relu` attribute, when set, lowers to the '.relu' variant
    of the cvt instruction.
    
    The optional scaling-factors for each of the inputs are provided through 
    the operand `scaleFactor` as a packed i16 type. Only `ue8m0` is supported 
    as the type of the scale-factor currently.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];

  let results = (outs VectorOfLengthAndType<[2], [BF16]>:$dst);
  let arguments = (ins VectorOfLengthAndType<[2], [I8]>:$src,
      Optional<I16>:$scaleFactor,
      DefaultValuedAttr<SaturationModeAttr, "SaturationMode::NONE">:$sat,
      DefaultValuedAttr<BoolAttr, "false">:$relu);
  let assemblyFormat =
      "$src (`,` $scaleFactor^)? attr-dict `:` type($src) `->` type($dst)";
}

//===----------------------------------------------------------------------===//
// NVVM Stochastic Rounding Conversion Ops
//===----------------------------------------------------------------------===//

// Base class for conversions from F32x2 to FPx2 formats
// (F16x2, BF16x2)
class NVVM_ConvertF32x2ToFPx2OpBase<string dstFormat, string mnemonic, Type dstType> :
  NVVM_Op<mnemonic, [Pure]>,
  Results<(outs dstType:$dst)>,
  Arguments<(ins F32:$src_hi, F32:$src_lo,
                 Optional<I32>:$random_bits,
                 DefaultValuedAttr<FPRoundingModeAttr, "FPRoundingMode::NONE">:$rnd,
                 DefaultValuedAttr<SaturationModeAttr, "SaturationMode::NONE">:$sat,
                 DefaultValuedAttr<BoolAttr, "false">:$relu)> {
  let summary = "Convert two F32 values to packed " # !tolower(dstFormat) # ".";
  let description = [{
    Converts two F32 values to packed }] # !tolower(dstFormat) # [{ format with 
    the specified rounding mode. The `src_hi` and `src_lo` parameters 
    correspond to operands `a` and `b` in the PTX ISA, respectively.
    
    The `random_bits` parameter is required for stochastic rounding and 
    provides the [random bits](}] #
    !if(!eq(dstFormat, "F16x2"),
    "https://docs.nvidia.com/cuda/parallel-thread-execution/#cvt-rs-rbits-layout-f16",
    "https://docs.nvidia.com/cuda/parallel-thread-execution/#cvt-rs-rbits-layout-bf16") #
    [{) to be used for the conversion.

    The `relu` attribute clamps negative results to 0.

    The `sat` attribute determines saturation behavior.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];
  
  let assemblyFormat = "$src_hi `,` $src_lo (`,` $random_bits^)? attr-dict `:` type($dst)";

  let hasVerifier = 1;
  
  let extraClassDeclaration = [{
    static NVVM::IDArgPair
    getIntrinsicIDAndArgs(
      NVVM::ConvertF32x2To}] # dstFormat # [{Op &op, 
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];
  
  string llvmBuilder = [{
    auto [intId, args] = mlir::NVVM::ConvertF32x2To}] # dstFormat # 
    [{Op::getIntrinsicIDAndArgs(op, moduleTranslation, builder);
    $dst = createIntrinsicCall(builder, intId, args);
  }];
}

// F32x2 -> F16x2
def NVVM_ConvertF32x2ToF16x2Op : NVVM_ConvertF32x2ToFPx2OpBase<"F16x2", "convert.f32x2.to.f16x2", VectorOfLengthAndType<[2], [F16]>>;

// F32x2 -> BF16x2
def NVVM_ConvertF32x2ToBF16x2Op : NVVM_ConvertF32x2ToFPx2OpBase<"BF16x2", "convert.f32x2.to.bf16x2", VectorOfLengthAndType<[2], [BF16]>>;

// Base class for stochastic rounding conversions from F32x4 to FPx4 formats
// (E4M3x4, E5M2x4, E2M3x4, E3M2x4, E2M1x4)
// These operations always use RS (stochastic rounding) mode with SATFINITE saturation.
class NVVM_ConvertF32x4ToFPx4OpBase<string dstFormat, string mnemonic, Type dstType> :
  NVVM_Op<mnemonic, [Pure, NVVMRequiresSMa<[100, 103]>]>,
  Results<(outs dstType:$dst)>,
  Arguments<(ins VectorOfLengthAndType<[4], [F32]>:$src, I32:$rbits,
                 DefaultValuedAttr<BoolAttr, "false">:$relu,
                 TypeAttr:$dstTy)> {
  let summary = "Convert vector<4xf32> to packed " # dstFormat # " with stochastic rounding (.rs) and satfinite";
  let description = [{
    Converts a vector<4xf32> to packed }] # dstFormat # [{ format using 
    stochastic rounding (.rs) mode with SATFINITE saturation. Randomness is 
    provided by the `rbits` parameter. The `dstTy` attribute specifies the 
    target floating-point format. The `relu` attribute clamps negative results to 0.
    
    Note: These operations always use RS rounding mode and SATFINITE saturation mode.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cvt)
  }];
  
  let assemblyFormat = "$src `,` $rbits attr-dict `:` type($src) `->` type($dst) `(` $dstTy `)`";

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    llvm::Intrinsic::ID getIntrinsicID();
  }];
  
  string llvmBuilder = [{
    auto intId = op.getIntrinsicID();
    $dst = createIntrinsicCall(builder, intId, {$src, $rbits});
  }];
}

// F32x4 -> F8x4 with stochastic rounding (supports E4M3FN, E5M2)
def NVVM_ConvertF32x4ToF8x4Op : NVVM_ConvertF32x4ToFPx4OpBase<"f8x4", "convert.f32x4.to.f8x4", VectorOfLengthAndType<[4], [I8]>>;

// F32x4 -> F6x4 with stochastic rounding (supports E2M3FN, E3M2FN)
def NVVM_ConvertF32x4ToF6x4Op : NVVM_ConvertF32x4ToFPx4OpBase<"f6x4", "convert.f32x4.to.f6x4", VectorOfLengthAndType<[4], [I8]>>;

// F32x4 -> F4x4 with stochastic rounding (supports E2M1FN)
def NVVM_ConvertF32x4ToF4x4Op : NVVM_ConvertF32x4ToFPx4OpBase<"f4x4", "convert.f32x4.to.f4x4", I16>;

//===----------------------------------------------------------------------===//
// NVVM MMA Ops
//===----------------------------------------------------------------------===//
/// Helpers to instantiate different version of wmma intrinsics.
/// This matches the hierarchy used in IntrinsicsNVVM.td to define all the
/// combinations of the intrinsics.
class GEOM<int M, int N, int K> {
  int m = M;
  int n = N;
  int k = K;
}

/// Class containing information about valid mma matrix types.
class WMMA_REGS<GEOM Geom, string Frag, string PtxEltType> {
  int m = Geom.m;
  int n = Geom.n;
  int k = Geom.k;
  string geom = "m"#Geom.m#"n"#Geom.n#"k"#Geom.k;
  string frag = Frag;
  string ptx_elt_type = PtxEltType;
  string gft = geom#":"#Frag#":"#ptx_elt_type;
}

//// Generate enum value of the mma.load/mma.store intrinsic.
class WMMA_NAME_LDST<string Op, WMMA_REGS Frag, string Layout, int WithStride> {
  string id =   "llvm::Intrinsic::nvvm_wmma"
                # "_" # Frag.geom
                # "_" # Op
                # "_" # Frag.frag
                # "_" # Frag.ptx_elt_type
                # "_" # Layout
                # !if(WithStride, "_stride", "");
}

/// Generate the signature part of the mma intrinsic name.
class MMA_SIGNATURE<WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  list<WMMA_REGS> id_frags = !cond(
     // FP8/F8F6F4 ops are identified by A,B inputs & accomulator & result type.
     !or(!eq(A.ptx_elt_type, "e4m3"),
         !eq(A.ptx_elt_type, "e5m2"),
         !eq(A.ptx_elt_type, "e3m2"),
         !eq(A.ptx_elt_type, "e2m3"),
         !eq(A.ptx_elt_type, "e2m1")): [D, A, B, C],
     // FP16 ops are identified by accumulator & result type.
     !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));
}

/// Generate enum value of the wmma.mma intrinsic.
class WMMA_NAME<string Op, string ALayout, string BLayout, WMMA_REGS A,
  WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  string signature = MMA_SIGNATURE<A, B, C, D>.ret;
  string id =   "llvm::Intrinsic::nvvm_wmma"
                # "_" # A.geom
                # "_" # Op
                # "_" # ALayout
                # "_" # BLayout
                # signature;
}

// Generates list of 4-tuples of WMMA_REGS representing a valid MMA op.
//   Geom: list of supported geometries.
//   TypeN: PTX type of the corresponding fragment's element.
//   TypeB and TypeD may be empty if it must match that of TypeA or TypeC.
class MMA_OPS<list<GEOM> Geom, list<string> TypeA, list<string> TypeB,
            list<string> TypeC, list<string> TypeD> {
  list<list<WMMA_REGS>> ret =
     !foldl([]<list<WMMA_REGS>>, Geom, t1, geom, !listconcat(t1,
     !foldl([]<list<WMMA_REGS>>, TypeA, t2, type_a, !listconcat(t2,
     !foldl([]<list<WMMA_REGS>>, !if(!size(TypeB), TypeB, [type_a]), t3, type_b, !listconcat(t3,
     !foldl([]<list<WMMA_REGS>>, TypeC, t4, type_c, !listconcat(t4,
     !foldl([]<list<WMMA_REGS>>, !if(!size(TypeD), TypeD, [type_c]), t5, type_d, !listconcat(t5,
            [[WMMA_REGS<geom, "a", type_a>,
              WMMA_REGS<geom, "b", type_b>,
              WMMA_REGS<geom, "c", type_c>,
              WMMA_REGS<geom, "d", type_d>]]))))))))));
   // Debugging aid for readable representation of the list above.
   list<list<string>> ops = !foreach(x, ret, [x[0].gft, x[1].gft, x[2].gft, x[3].gft]);
}

/// Creates a list of combinations of load/store operations supported.
class MMA_LDST_OPS<list<GEOM> Geom, list<string> Frags, list<string> Types> {
  list<WMMA_REGS> ret =
     !foldl([]<WMMA_REGS>, Geom, t1, geom, !listconcat(t1,
     !foldl([]<WMMA_REGS>, Frags, t2, frag, !listconcat(t2,
     !foldl([]<WMMA_REGS>, Types, t3, type, !listconcat(t3,
            [WMMA_REGS<geom, frag, type>]))))));
   // Debugging aid for readable representation of the list above.
   list<string> ops = !foreach(x, ret, x.gft);
}

// Creates list of valid combinations of fragments. This is a subset of what
// llvm supports and can be extended as needed.
class NVVM_MMA_OPS {
  // "wmma" operations
  list<list<WMMA_REGS>> fp64_wmma_ops = MMA_OPS<
            [GEOM<8, 8, 4>],
            ["f64"], [], ["f64"], []>.ret;
  list<list<WMMA_REGS>> tf32_wmma_ops = MMA_OPS<
            [GEOM<16, 16, 8>],
            ["tf32"], [], ["f32"], []>.ret;
  list<list<WMMA_REGS>> fp_wmma_ops = MMA_OPS<
            [GEOM<16, 16, 16>, GEOM<32, 8, 16>, GEOM<8, 32, 16>],
            ["f16"], [], ["f16", "f32"], []>.ret;
  list<list<WMMA_REGS>> i8_wmma_ops = MMA_OPS<
            [GEOM<16, 16, 16>, GEOM<32, 8, 16>, GEOM<8, 32, 16>],
            ["s8","u8"], [], ["s32"], []>.ret;
  list<list<WMMA_REGS>> all_wmma_ops = !listconcat(
            fp64_wmma_ops,
            tf32_wmma_ops,
            fp_wmma_ops,
            i8_wmma_ops);

  list<WMMA_REGS> ldst_ab_ops = MMA_LDST_OPS<
            [GEOM<16, 16, 16>, GEOM<32, 8, 16>, GEOM<8, 32, 16>],
            ["a", "b"], ["f16","s8","u8"]>.ret;
  list<WMMA_REGS> ldst_cd_ops = MMA_LDST_OPS<
            [GEOM<16, 16, 16>, GEOM<32, 8, 16>, GEOM<8, 32, 16>],
            ["c", "d"], ["f16", "f32","s32"]>.ret;
  list<WMMA_REGS> ldst_tf32_ab_ops = MMA_LDST_OPS<
            [GEOM<16, 16, 8>],
            ["a", "b"], ["tf32"]>.ret;
  list<WMMA_REGS> ldst_tf32_cd_ops = MMA_LDST_OPS<
            [GEOM<16, 16, 8>],
            ["c", "d"], ["f32"]>.ret;
  list<WMMA_REGS> ldst_f64_ab_ops = MMA_LDST_OPS<
            [GEOM<8, 8, 4>],
            ["a", "b"], ["f64"]>.ret;
  list<WMMA_REGS> ldst_f64_cd_ops = MMA_LDST_OPS<
            [GEOM<8, 8, 4>],
            ["c", "d"], ["f64"]>.ret;
  list<WMMA_REGS> all_ldst_ops = !listconcat(ldst_ab_ops, ldst_cd_ops,
                                             ldst_tf32_ab_ops,
                                             ldst_tf32_cd_ops, 
                                             ldst_f64_ab_ops,
                                             ldst_f64_cd_ops);
  // Separate A/B/C fragments (loads) from D (stores).
  list<WMMA_REGS> all_ld_ops = !filter(op, all_ldst_ops, !ne(op.frag, "d"));
  list<WMMA_REGS> all_st_ops = !filter(op, all_ldst_ops, !eq(op.frag, "d"));

  // "mma_sync" operations
  list<list<WMMA_REGS>> tf32_mma_ops = MMA_OPS<
            [GEOM<16,8,4>, GEOM<16,8,8>],
            ["tf32"], [], ["f32"], []>.ret;
  list<list<WMMA_REGS>> bf16_mma_ops = MMA_OPS<
            [GEOM<16,8,16>, GEOM<16,8,8>],
            ["bf16"], [], ["f32"], []>.ret;
  list<list<WMMA_REGS>> f64_mma_ops = MMA_OPS<
            [GEOM<8,8,4>],
            ["f64"], [], ["f64"], []>.ret;
  list<list<WMMA_REGS>> fp_mma_ops = MMA_OPS<
            [GEOM<8,8,4>, GEOM<16,8,8>, GEOM<16,8,16>],
            ["f16"], [], ["f16", "f32"], ["f16", "f32"]>.ret;
  list<list<WMMA_REGS>> int_mma_ops = MMA_OPS<
            [GEOM<8,8,16>, GEOM<16,8,16>, GEOM<16,8,32>],
            ["s8", "u8"], ["s8", "u8"], ["s32"], []>.ret;
  list<list<WMMA_REGS>> subint_mma_ops = MMA_OPS<
            [GEOM<8,8,32>, GEOM<16,8,32>, GEOM<16,8,64>],
            ["s4", "u4"], ["s4", "u4"], ["s32"], []>.ret;
  list<list<WMMA_REGS>> bit_mma_ops = MMA_OPS<
            [GEOM<8,8,128>, GEOM<16,8,128>, GEOM<16,8,256>],
            ["b1"], [], ["s32"], []>.ret;
  list<list<WMMA_REGS>> all_mma_sync_ops = !listconcat(
            tf32_mma_ops, bf16_mma_ops, f64_mma_ops,
            fp_mma_ops, int_mma_ops, subint_mma_ops, bit_mma_ops);

  list<list<WMMA_REGS>> bf16_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,16>, GEOM<16,8,32>],
            ["bf16"], [], ["f32"], []>.ret;
  list<list<WMMA_REGS>> tf32_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,8>, GEOM<16,8,16>],
            ["tf32"], [], ["f32"], []>.ret;
  list<list<WMMA_REGS>> fp_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,16>, GEOM<16,8,32>],
            ["f16"], [], ["f16", "f32"], ["f16", "f32"]>.ret;
  list<list<WMMA_REGS>> fp8_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,64>],
            ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],
            ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],
            ["f16", "f32"], ["f16", "f32"]>.ret;
  list<list<WMMA_REGS>> subint_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,64>, GEOM<16,8,128>],
            ["s4", "u4"], ["s4", "u4"], ["s32"], []>.ret;
  list<list<WMMA_REGS>> int_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,32>, GEOM<16,8,64>],
            ["s8", "u8"], ["s8", "u8"], ["s32"], []>.ret;
  list<list<WMMA_REGS>> all_mma_sp_sync_ops = !listconcat(
            bf16_mma_sp_ops, tf32_mma_sp_ops, fp_mma_sp_ops, fp8_mma_sp_ops,
            subint_mma_sp_ops, int_mma_sp_ops);

  // Block scale MMA operations (dense)
  list<list<WMMA_REGS>> mxf4_mma_ops = MMA_OPS<
            [GEOM<16,8,64>],
            ["e2m1"], ["e2m1"], ["f32"], []>.ret;
  list<list<WMMA_REGS>> mxf8f6f4_mma_ops = MMA_OPS<
            [GEOM<16,8,32>],
            ["e2m1", "e2m3", "e3m2", "e5m2", "e4m3"],
            ["e2m1", "e2m3", "e3m2", "e5m2", "e4m3"],
            ["f32"], []>.ret;
  list<list<WMMA_REGS>> all_mma_block_scale_ops = !listconcat(
            mxf4_mma_ops, mxf8f6f4_mma_ops);

  // Block scale sparse MMA operations
  list<list<WMMA_REGS>> mxf4xx_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,128>],
            ["e2m1"], ["e2m1"], ["f32"], []>.ret;
  list<list<WMMA_REGS>> mxf8f6f4_mma_sp_ops = MMA_OPS<
            [GEOM<16,8,64>],
            ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],
            ["e4m3", "e5m2", "e3m2", "e2m3", "e2m1"],
            ["f32"], []>.ret;
  list<list<WMMA_REGS>> all_mma_sp_block_scale_ops = !listconcat(
            mxf4xx_mma_sp_ops, mxf8f6f4_mma_sp_ops);

}

def NVVM_MMA_OPS : NVVM_MMA_OPS;

/// Helper to create the mapping between the configuration and the store
/// intrinsic enum value.
class MMA_ST_INTR<string op> {
  list<list<string>> cond0 = !foreach(frag, NVVM_MMA_OPS.all_st_ops,
                                !foreach(layout, ["row", "col"],
  "if (layout == \"" # layout #  "\" && m == " # frag.m # " &&"
  "    n == " #frag.n # " && k == " # frag.k # " && \"" #
       frag.ptx_elt_type # "\" == eltype)"
  "  return " #WMMA_NAME_LDST<op, frag, layout, 1>.id #";"));
  string id = !foldl("",
                !foldl([""], cond0, acc, el, !listconcat(acc, el)),
                acc1, el1, acc1 # "\n" # el1);
}

/// Helper to map a mxk shape to a supported mxnxk matrix type. This will return
/// the n value of the supported configuration.
class MMA_ST_INFER_N<list<WMMA_REGS> ldst> {
  list<string> cond = !foreach(frag, ldst,
  "if (m == " # frag.m # " && k == " #frag.k # " && \"" #
       frag.ptx_elt_type # "\" == eltype)"
  "  return "# frag.n #";");
  string id = !foldl("", cond, acc, el, acc # "\n" # el);
}

/// Helper to map a kxn shape to a supported mxnxk matrix type. This will return
/// the m value of the supported configuration.
class MMA_ST_INFER_M<list<WMMA_REGS> ldst> {
  list<string> cond = !foreach(frag, ldst,
  "if (n == " # frag.n # " && k == " #frag.k # " && \"" #
       frag.ptx_elt_type # "\" == eltype)"
  "  return "# frag.m #";");
  string id = !foldl("", cond, acc, el, acc # "\n" # el);
}

/// Helper to map a mxn shape to a supported mxnxk matrix type. This will return
/// the k value of the supported configuration.
class MMA_ST_INFER_K<list<WMMA_REGS> ldst> {
  list<string> cond = !foreach(frag, ldst,
  "if (m == " # frag.m # " && n == " #frag.n # " && \"" #
       frag.ptx_elt_type # "\" == eltype)"
  "  return "# frag.k #";");
  string id = !foldl("", cond, acc, el, acc # "\n" # el);
}

/// Helper to create the mapping between the configuration and the load
/// intrinsic enum value.
class MMA_LD_INTR<string op> {
  list<list<string>> cond0 = !foreach(frag, NVVM_MMA_OPS.all_ld_ops,
                                !foreach(layout, ["row", "col"],
  "if (layout == \"" # layout #  "\" && m == " # frag.m # " &&"
  "    n == " #frag.n # " && k == " # frag.k # " && \"" #
       frag.ptx_elt_type # "\" == eltype && frag == \""#frag.frag#"\")"
  "  return "# WMMA_NAME_LDST<op, frag, layout, 1>.id #";"));
  string id = !foldl("",
                !foldl([""], cond0, acc, el, !listconcat(acc, el)),
                acc1, el1, acc1 # "\n" # el1);
}

/// Helper to create the mapping between the configuration and the wmma.mma
/// intrinsic enum value.
class MMA_MMA_INTR<string opName> {
  list<list<list<string>>> cond0 =
    !foreach(op, NVVM_MMA_OPS.all_wmma_ops,
      !foreach(layoutA, ["row", "col"],
        !foreach(layoutB, ["row", "col"],
  "if (layoutA == \"" # layoutA #  "\" && layoutB == \"" # layoutB #  "\" && "
  "    m == " # op[0].m # " && n == " #op[0].n # " && k == " # op[0].k #
  "    && \"" # op[0].ptx_elt_type # "\" == eltypeA && \""
   # op[3].ptx_elt_type # "\" == eltypeB)"
  "  return " #
       WMMA_NAME<opName, layoutA, layoutB, op[0], op[1], op[2], op[3]>.id # ";")));
  list<string> f = !foldl([""],
                     !foldl([[""]], cond0, acc, el, !listconcat(acc, el)),
                          acc1, el1, !listconcat(acc1, el1));
  string id = !foldl("", f, acc, el, acc # "\n" # el);
}

/// Enum attribute for binary (b1) MMA operation type
def MMAB1OpNone : I32EnumAttrCase<"none", 0>;
def MMAB1OpXorPopc : I32EnumAttrCase<"xor_popc", 1>;
def MMAB1OpAndPopc : I32EnumAttrCase<"and_popc", 2>;
def MMAB1Op : I32EnumAttr<"MMAB1Op", "MMA binary operations",
  [MMAB1OpNone, MMAB1OpXorPopc, MMAB1OpAndPopc]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMAB1OpAttr : EnumAttr<NVVM_Dialect, MMAB1Op, "mma_b1op"> {
  let assemblyFormat = "`<` $value `>`";
}

/// Enum attribute type for the overflow behavior of MMA integer operations
def MMAIntOverflowWrap : I32EnumAttrCase<"wrapped", 0>;
def MMAIntOverflowSat : I32EnumAttrCase<"satfinite", 1>;
def MMAIntOverflow : I32EnumAttr<"MMAIntOverflow", "MMA overflow options",
  [MMAIntOverflowSat, MMAIntOverflowWrap]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMAIntOverflowAttr : EnumAttr<NVVM_Dialect, MMAIntOverflow, "mma_int_overflow"> {
  let assemblyFormat = "`<` $value `>`";
}
/// MMA kind types (for mixed-precision FP8 operations)
def MMAKindF8F6F4 : I32EnumAttrCase<"f8f6f4", 0>;
def MMAKind : I32EnumAttr<"MMAKind", "MMA operation kind",
  [MMAKindF8F6F4]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMAKindAttr : EnumAttr<NVVM_Dialect, MMAKind, "mma_kind"> {
  let assemblyFormat = "`<` $value `>`";
}

/// Attribute to hold the MMA shape
def NVVM_MMAShapeAttr : NVVM_Attr<"MMAShape", "shape"> {
  let summary = "Attribute for MMA operation shape.";
  let parameters = (ins "int":$m, "int":$n, "int":$k);
  let assemblyFormat = "`<` struct(params) `>`";
}

// 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.
  string layout = layout_a # ":" # layout_b;
  string a_type = frags[0].ptx_elt_type;
  string b_type = frags[1].ptx_elt_type;
  string c_type = frags[2].ptx_elt_type;
  string d_type = frags[3].ptx_elt_type;
  string geom = frags[0].geom;

  // gcd is a shortcut used to identify instructions that depend on
  // geom+frag_c+frag_d.
  string gcd = geom # ":" # c_type # d_type;
  bit ret = !cond(

    // Limit satf to valid types
    !and(!eq(satf, 1),
         !ne(a_type, "s8"),
         !ne(a_type, "u8"),
         !ne(a_type, "s4"),
         !ne(a_type, "u4")): false,

    // m8n8k4 has no C=f32 D=f16 variant.
    !eq(gcd, "m8n8k4:f32f16"): false,

    // only m8n8k4 for f16 does not require row:col layout
    !and(!ne(layout, "row:col"),
         !or(!ne(geom, "m8n8k4"),
             !ne(a_type, "f16"))) : false,

    // m16n8k8 requires A and B to be the same type and C and D to be the same
    // type.
    !and(!eq(geom, "m16n8k8"),
         !or(!ne(a_type, b_type),
             !ne(c_type, d_type))): false,

    // m16n8k16/m16n8k32 requires C and D to be the same type
    !and(!or(!eq(geom, "m16n8k16"),
             !eq(geom, "m16n8k32")),
         !ne(c_type, d_type)): false,

    // All other are OK.
    true: true
  );
}

// Returns a list of operation suffixes corresponding to possible b1
// multiply-and-accumulate operations for all fragments which have a
// b1 type. For all other fragments, the list returned holds a list
// containing the empty string.
class NVVM_MMA_B1OPS<list<WMMA_REGS> frags> {
  list<string> ret = !cond(
    !eq(frags[0].ptx_elt_type, "b1") : ["xor_popc", "and_popc"],
    true: [""]
  );
}

/// Generate enum value of the mma.sync intrinsic.
class MMA_SYNC_NAME<string ALayout, string BLayout, string b1op, int Satfinite,
               WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  string signature = MMA_SIGNATURE<A, B, C, D>.ret;
  string id = "llvm::Intrinsic::nvvm_mma"
                # !if(!ne(b1op, ""), "_" # b1op, "")
                # "_" # A.geom
                # "_" # ALayout
                # "_" # BLayout
                # !if(Satfinite, "_satfinite", "")
                # signature;
}

/// Helper to create the mapping between the configuration and the mma.sync
/// intrinsic enum value.
class MMA_SYNC_INTR {
  list<list<list<list<list<string>>>>> cond0 =
    !foreach(op, NVVM_MMA_OPS.all_mma_sync_ops,
      !foreach(layoutA, ["row", "col"],
        !foreach(layoutB, ["row", "col"],
          !foreach (sat, [0, 1],
            !foreach (b1op, NVVM_MMA_B1OPS<op>.ret,
              !if(NVVM_MMA_SUPPORTED<[op[0], op[1], op[2], op[3]],
                                     layoutA, layoutB, sat>.ret,
      "if (layoutA == \"" # layoutA #  "\" && layoutB == \"" # layoutB #  "\" && "
      "    m == " # op[0].m # " && n == " # op[0].n # " && k == " # op[0].k #
      "    && \"" # op[0].ptx_elt_type # "\" == eltypeA && \""
       # op[1].ptx_elt_type # "\" == eltypeB && "
       # " \"" # op[2].ptx_elt_type # "\" == eltypeC && "
       # " \"" # op[3].ptx_elt_type # "\" == eltypeD "
       # " && (sat.has_value()  ? " # sat # " == static_cast<int>(*sat) : true)"
       # !if(!ne(b1op, ""), " && (b1Op.has_value() ? MMAB1Op::" # b1op # " == *b1Op : true)", "") # ")\n"
       # "  return " #
       MMA_SYNC_NAME<layoutA, layoutB, b1op, sat, op[0], op[1], op[2], op[3]>.id # ";",
          "") // if supported
          ) // b1op
        ) // sat
      ) // layoutB
    ) // layoutA
  ); // all_mma_sync_ops
  list<list<list<string>>> f1 = !foldl([[[""]]],
                                  !foldl([[[[""]]]], cond0, acc, el,
                                      !listconcat(acc, el)),
                                    acc1, el1, !listconcat(acc1, el1));
  list<list<string>> f2 = !foldl([[""]], f1, acc1, el1, !listconcat(acc1, el1));
  list<string> f3 = !foldl([""], f2, acc, el, !listconcat(acc, el));
  string id = !foldl("", f3, acc, el, acc # "\n" # el);
}

def MMALayoutRow : I32EnumAttrCase<"row", 0>;
def MMALayoutCol : I32EnumAttrCase<"col", 1>;

/// Enum attribute of the different matrix layout.
def MMALayout : I32EnumAttr<"MMALayout", "NVVM MMA layout",
  [MMALayoutRow, MMALayoutCol]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMALayoutAttr : EnumAttr<NVVM_Dialect, MMALayout, "mma_layout"> {
  let assemblyFormat = "`<` $value `>`";
}

/// Enum attribute of the different PTX element types used for MMA operands.
def MMATypeF16  : I32EnumAttrCase<"f16", 0>;
def MMATypeF32  : I32EnumAttrCase<"f32", 1>;
def MMATypeTF32 : I32EnumAttrCase<"tf32", 2>;
def MMATypeU8 : I32EnumAttrCase<"u8", 3>;
def MMATypeS8 : I32EnumAttrCase<"s8", 4>;
def MMATypeS32 : I32EnumAttrCase<"s32", 5>;
def MMATypeB1 : I32EnumAttrCase<"b1", 6>;
def MMATypeU4 : I32EnumAttrCase<"u4", 7>;
def MMATypeS4 : I32EnumAttrCase<"s4", 8>;
def MMATypeBF16 : I32EnumAttrCase<"bf16", 9>;
def MMATypeF64 : I32EnumAttrCase<"f64", 10>;
def MMATypeE4M3 : I32EnumAttrCase<"e4m3", 11>;
def MMATypeE5M2 : I32EnumAttrCase<"e5m2", 12>;
def MMATypeE3M2 : I32EnumAttrCase<"e3m2", 13>;
def MMATypeE2M3 : I32EnumAttrCase<"e2m3", 14>;
def MMATypeE2M1 : I32EnumAttrCase<"e2m1", 15>;

def MMATypes : I32EnumAttr<"MMATypes", "NVVM MMA types",
  [MMATypeF16, MMATypeF32, MMATypeTF32,
  MMATypeBF16, MMATypeS8, MMATypeU8,
  MMATypeS32, MMATypeS4, MMATypeU4,
  MMATypeB1, MMATypeF64,
  MMATypeE4M3, MMATypeE5M2, MMATypeE3M2, MMATypeE2M3, MMATypeE2M1]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMATypesAttr : EnumAttr<NVVM_Dialect, MMATypes, "mma_type"> {
  let assemblyFormat = "`<` $value `>`";
}

def MMAFragA : I32EnumAttrCase<"a", 0>;
def MMAFragB : I32EnumAttrCase<"b", 1>;
def MMAFragC : I32EnumAttrCase<"c", 2>;

/// Enum attribute of the different frag types.
def MMAFrag: I32EnumAttr<"MMAFrag", "NVVM MMA frag type",
  [MMAFragA, MMAFragB, MMAFragC]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def MMAFragAttr : EnumAttr<NVVM_Dialect, MMAFrag, "mma_frag"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_WMMALoadOp: NVVM_Op<"wmma.load">,
  Results<(outs AnyTypeOf<[LLVM_AnyStruct, F64]>:$res)>,
  Arguments<(ins LLVM_AnyPointer: $ptr, I32: $stride, I32Attr:$m,
             I32Attr:$n, I32Attr:$k, MMALayoutAttr:$layout,
             MMATypesAttr:$eltype, MMAFragAttr:$frag)> {

  let summary = "Warp synchronous matrix load";

  // Since LLVM intrinsic IDs are enum that cannot be dynamically generated in
  // C++ we instanciate a function in tablegen to map the valide configuration
  // to the corresponsding intrinsic ID.
  // Because we want a single source of truth, this mean the source of truth
  // about valid combinations needs to be in tablgen, therefore we generate
  // extra helpers to query valid configurations based on the shapes of
  // load/store operations.
  let extraClassDeclaration =
    "static llvm::Intrinsic::ID getIntrinsicID("
    "int m, int n, int k, mlir::NVVM::MMALayout layoutEnum,"
    "mlir::NVVM::MMATypes eltypeEnum,mlir::NVVM::MMAFrag fragEnum) {"
    "llvm::StringRef layout = stringifyEnum(layoutEnum);"
    "llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    "llvm::StringRef frag = stringifyEnum(fragEnum);"
    #MMA_LD_INTR<"load">.id# "\n"
    "return 0;"
    "}\n"
    "/// Helpers to find valid n dimension based on mxk load shape.\n"
    "static int inferNDimension(int m, int k, mlir::NVVM::MMATypes eltypeEnum) {"
    "  llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    #MMA_ST_INFER_N<!filter(op, NVVM_MMA_OPS.all_ld_ops, !eq(op.frag, "a"))>.id# "\n"
    "return 0;"
    "}\n"
    "/// Helpers to find valid m dimension based on kxn load shape.\n"
    "static int inferMDimension(int k, int n, mlir::NVVM::MMATypes eltypeEnum) {"
    "  llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    #MMA_ST_INFER_M<!filter(op, NVVM_MMA_OPS.all_ld_ops, !eq(op.frag, "b"))>.id# "\n"
    "return 0;"
    "}\n"
    "/// Helpers to find valid k dimension based on mxn load shape.\n"
    "static int inferKDimension(int m, int n, mlir::NVVM::MMATypes eltypeEnum) {"
    "  llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    #MMA_ST_INFER_K<!filter(op, NVVM_MMA_OPS.all_ld_ops, !eq(op.frag, "c"))>.id# "\n"
    "return 0;"
    "}\n";


  string llvmBuilder = [{
      auto operands = moduleTranslation.lookupValues(opInst.getOperands());
      auto intId = mlir::NVVM::WMMALoadOp::getIntrinsicID(
        $m, $n, $k, $layout, $eltype, $frag);
      $res = createIntrinsicCall(builder, intId, operands, {operands[0]->getType()});
  }];

  string baseDescription = [{
    The `nvvm.wmma.load` operation loads a matrix collectively using all the
    threads in a warp.

    The operation takes two arguments, the address from where the matrix
    elements are to be loaded from and a stride. The stride argument
    represents the leading dimension of the source matrix. The address and
    the stride are required to be the same across all threads in the warp.
    Each thread in a warp holds a certain number of elements. The Op returns
    a LLVMStruct which holds the elements of the matrix held by this thread.

    This op is meant to be used along with `nvvm.wmma.store` and
    `nvvm.wmma.mma`.

    Example:

    ```mlir
    %2 = nvvm.wmma.load %0, %1
      {eltype = "f16", frag = "a", k = 16 : i32, layout = "row", m = 16 : i32, n = 16 : i32}
      : (!llvm.ptr<3>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
    ```
    }];

  let assemblyFormat = "$ptr `,` $stride attr-dict `:` functional-type($ptr, $res)";
  let hasVerifier = 1;
}

def NVVM_WMMAStoreOp : NVVM_Op<"wmma.store">,
  Arguments<(ins LLVM_AnyPointer: $ptr,
             I32Attr:$m, I32Attr:$n, I32Attr:$k, MMALayoutAttr:$layout,
             MMATypesAttr:$eltype, Variadic<LLVM_Type>:$args, I32: $stride)>{
  let summary = "Warp synchronous matrix store";

  let extraClassDeclaration =
    "static llvm::Intrinsic::ID getIntrinsicID("
    "int m, int n, int k, mlir::NVVM::MMALayout layoutEnum,"
    "mlir::NVVM::MMATypes eltypeEnum) {"
    "  llvm::StringRef layout = stringifyEnum(layoutEnum);"
    "  llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    #MMA_ST_INTR<"store">.id# "\n"
    "return 0;"
    "}\n"
    "/// Helpers to find valid k dimension based on mxn store shape.\n"
    "static int inferKDimension(int m, int n, mlir::NVVM::MMATypes eltypeEnum) {"
    "  llvm::StringRef eltype = stringifyEnum(eltypeEnum);"
    #MMA_ST_INFER_K<NVVM_MMA_OPS.all_st_ops>.id#  "\n"
    "return 0;"
    "}";

  string llvmBuilder = [{
      auto operands = moduleTranslation.lookupValues(opInst.getOperands());
      auto intId =
        mlir::NVVM::WMMAStoreOp::getIntrinsicID($m, $n, $k, $layout, $eltype);
      createIntrinsicCall(builder, intId, operands, {operands[0]->getType()});
  }];

  string baseDescription = [{
    The `nvvm.wmma.store` operation stores a matrix collectively using
    all the threads in a warp.

    The operation takes as arguments the address to where the matrix elements are
    to be stored, a stride and the elements to store, held by the current thread.
    The stride argument represents the leading dimension of the destination matrix.
    The address and the stride are required to be the same across all threads in the
    warp.

    This op is meant to be used along with `nvvm.wmma.m16n16k16.load` and
    `nvvm.wmma.m16n16k16.mma`.

    Example:

    ```mlir
    nvvm.wmma.store %0, %1, %2, %3, %4, %5
      {eltype = "f16", k = 16 : i32, layout = "row", m = 16 : i32, n = 16 : i32}
      : !llvm.ptr<3>, vector<2 x f16>, vector<2 x f16>, vector<2 x f16>, vector<2 x f16>
    ```
  }];

  let assemblyFormat = [{
    $ptr `,` $stride `,` $args attr-dict `:` qualified(type($ptr)) `,`
    type($args)
  }];
  let hasVerifier = 1;
}

// Base class for all the variants of WMMA mmaOps that may be defined.
def NVVM_WMMAMmaOp : NVVM_Op<"wmma.mma">,
  Results<(outs LLVM_AnyStruct:$res)>,
  Arguments<(ins I32Attr:$m, I32Attr:$n, I32Attr:$k, MMALayoutAttr:$layoutA,
             MMALayoutAttr:$layoutB, MMATypesAttr:$eltypeA,
             MMATypesAttr:$eltypeB, Variadic<LLVM_Type>:$args)>{
  let summary = "Warp synchronous matrix-multiply accumulate using tensor cores.";

  let extraClassDeclaration =
    "static llvm::Intrinsic::ID getIntrinsicID("
    "int m, int n, int k, mlir::NVVM::MMALayout layoutAEnum,"
    "mlir::NVVM::MMALayout layoutBEnum, mlir::NVVM::MMATypes eltypeAEnum,"
    "mlir::NVVM::MMATypes eltypeBEnum) {"
    "llvm::StringRef layoutA = stringifyEnum(layoutAEnum);"
    "llvm::StringRef layoutB = stringifyEnum(layoutBEnum);"
    "llvm::StringRef eltypeA = stringifyEnum(eltypeAEnum);"
    "llvm::StringRef eltypeB = stringifyEnum(eltypeBEnum);"
    #MMA_MMA_INTR<"mma">.id# "\n"
    "return 0;"
    "}";

  string llvmBuilder = [{
      auto operands = moduleTranslation.lookupValues(opInst.getOperands());
      auto intId = mlir::NVVM::WMMAMmaOp::getIntrinsicID(
        $m, $n, $k, $layoutA, $layoutB, $eltypeA, $eltypeB);
      $res = createIntrinsicCall(builder, intId, operands);
  }];

  string baseDescription = [{
    The `nvvm.wmma.mma` operation performs a matrix-multiply accumulate
    (mma) operation using all the threads in a warp.

    The operation performed is represented as `D = A * B + C`. The operation takes
    as arguments the elements of the matrices `A`, `B`, `C` and `D`, held by the
    current thread. The op returns a LLVM struct which holds a part of the result
    held by the current thread.

    This op is meant to be used along with `nvvm.wmma.load` and
    `nvvm.wmma.store`.

    Example:

    ```mlir
    %16 = nvvm.wmma.mma %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15
      {eltypeA = "tf32", eltypeB = "f32", k = 8 : i32, layoutA = "row", layoutB = "row", m = 16 : i32, n = 16 : i32}
      : (i32, i32, i32, i32, i32, i32, i32, i32, f32, f32, f32, f32, f32, f32, f32, f32)
      -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
    ```
  }];

  let assemblyFormat = "$args attr-dict `:` functional-type($args, $res)";
  let hasVerifier = 1;
}

def LdStMatrixShapeAttr : NVVM_Attr<"LdStMatrixShape", "ld_st_matrix_shape"> {
  let summary = "Matrix shape for ldmatrix and stmatrix";
  let parameters = (ins "int":$m, "int":$n);
  let assemblyFormat = "`<` struct(params) `>`";
}

def LdStMatrixEltTypeB16 : I32EnumAttrCase<"B16", 0, "b16">;
def LdStMatrixEltTypeB8 : I32EnumAttrCase<"B8", 1, "b8">;
def LdStMatrixEltTypeB8X16_B6X16_P32 : I32EnumAttrCase<"B8X16_B6X16_P32", 2, "b8x16.b6x16_p32">;
def LdStMatrixEltTypeB8X16_B4X16_P64 : I32EnumAttrCase<"B8X16_B4X16_P64", 3, "b8x16.b4x16_p64">;

def LdStMatrixEltType : I32EnumAttr<"LdStMatrixEltType", "Element type for ldmatrix and stmatrix",
  [LdStMatrixEltTypeB16, LdStMatrixEltTypeB8,
   LdStMatrixEltTypeB8X16_B6X16_P32, LdStMatrixEltTypeB8X16_B4X16_P64]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def LdStMatrixEltTypeAttr : EnumAttr<NVVM_Dialect, LdStMatrixEltType, "ld_st_matrix_elt_type"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_StMatrixOp: NVVM_Op<"stmatrix">,
  Arguments<(ins LLVM_PointerShared: $ptr, Variadic<I32>:$sources, MMALayoutAttr:$layout,
             LdStMatrixShapeAttr:$shape, LdStMatrixEltTypeAttr:$eltType)> {
  let summary = "cooperative matrix store";
  let description = [{
    Collectively store one or more matrices across all threads in a warp to the
    location indicated by the address operand $ptr in shared memory.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-store-instruction-stmatrix)
  }];
  string llvmBuilder = [{
      auto operands = moduleTranslation.lookupValues(opInst.getOperands());
      auto intId = getStMatrixIntrinsicId($layout, $sources.size(), $shape, $eltType);
      createIntrinsicCall(builder, intId, operands, operands[0]->getType());
  }];
  let assemblyFormat = "$ptr `,` $sources attr-dict `:` type(operands)";
  let hasVerifier = 1;
}

def NVVM_LdMatrixOp: NVVM_Op<"ldmatrix", [InferTypeOpAdaptor]>,
  Results<(outs AnyType:$res)>,
  Arguments<(ins LLVM_PointerShared:$ptr, I32Attr:$num,
                 MMALayoutAttr:$layout,
                 LdStMatrixShapeAttr:$shape,
                 LdStMatrixEltTypeAttr:$eltType)> {

  let summary = "cooperative matrix load";

  string llvmBuilder = [{
      auto operands = moduleTranslation.lookupValues(opInst.getOperands());
      auto intId = getLdMatrixIntrinsicId($layout, $num, $shape, $eltType);
      $res = createIntrinsicCall(builder, intId, operands, {operands[0]->getType()});
  }];

  string baseDescription = [{
    The `nvvm.ldmatrix` operation collectively loads one or more matrices across
    all threads in a warp from the location indicated by the address operand
    `ptr` from shared memory.

    The attribute `num` indicates how many 8x8 16-bit matrices are to be loaded.

    All the threads in the warp must execute the same ldmatrix operations.

    Each row of 8 elements needs to be consecutive in memory. Each lane of the
    warp contains the start address of a row of 8 elements laid out as below:

    ```
    num | lane 0--7    | Threads 8--15  | Threads 16--31
    1   | addr0--addr7 |                |
    2   | addr0--addr7 | addr8--addr15  |
    4   | addr0--addr7 | addr8--addr15  | addr16--addr31
    ```

    Example:
    ```mlir
    %l1 = nvvm.ldmatrix %ptr {num = 1 : i32, layout = #nvvm.mma_layout<row>} :
      (!llvm.ptr<3>) -> i32
    %l2 = nvvm.ldmatrix %ptr {num = 4 : i32, layout = #nvvm.mma_layout<row>} :
      (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
    ```
  }];

  let assemblyFormat = "$ptr attr-dict `:` functional-type($ptr, $res)";
  let hasVerifier = 1;
}

def NVVM_MovMatrixOp
    : NVVM_SingleResultIntrinsicOp<"movmatrix", [NVVMRequiresSM<75>], "$dst"> {
  let summary = "Warp-level matrix transpose";
  let description = [{
    Moves a row-major matrix across all threads in a warp, reading elements
    from source `$src`, and writing the transposed elements to destination
    `$dst`.

    The `shape` attribute indicates the dimensions of the matrix being
    transposed. Each matrix element holds 16-bit data as indicated by the
    `eltType` attribute.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-movmatrix-instruction)

    Example:
    ```mlir
    %dst = nvvm.movmatrix %src {shape = #nvvm.ld_st_matrix_shape<m = 8, n = 8>,
                                eltType = #nvvm.ld_st_matrix_elt_type<b16>} : i32
    ```
  }];

  let results = (outs I32:$dst);
  let arguments = (ins I32:$src, LdStMatrixShapeAttr:$shape,
      DefaultValuedAttr<MMALayoutAttr, "MMALayout::col">:$layout,
      LdStMatrixEltTypeAttr:$eltType);

  let assemblyFormat = "$src attr-dict `:` type($src)";
  let hasVerifier = 1;
}

def NVVM_MmaOp : NVVM_Op<"mma.sync", [AttrSizedOperandSegments]> {

  let summary = "cooperative matrix-multiply and accumulate";

  let description = [{
    The `nvvm.mma.sync` operation collectively performs the operation
    `D = matmul(A, B) + C` using all threads in a warp.

    All the threads in the warp must execute the same `mma.sync` operation.

    For each possible multiplicand PTX data type, there are one or more possible
    instruction shapes given as "mMnNkK". The below table describes the posssibilities
    as well as the types required for the operands. Note that the data type for
    C (the accumulator) and D (the result) can vary independently when there are
    multiple possibilities in the "C/D Type" column.

    When an optional attribute cannot be immediately inferred from the types of
    the operands and the result during parsing or validation, an error will be
    raised.

    `b1Op` is only relevant when the binary (b1) type is given to
    `multiplicandDataType`. It specifies how the multiply-and-acumulate is
    performed and is either `xor_popc` or `and_poc`. The default is `xor_popc`.

    `intOverflowBehavior` is only relevant when the `multiplicandType` attribute
    is one of `u8, s8, u4, s4`, this attribute describes how overflow is handled
    in the accumulator. When the attribute is `satfinite`, the accumulator values
    are clamped in the int32 range on overflow. This is the default behavior.
    Alternatively, accumulator behavior `wrapped` can also be specified, in
    which case overflow wraps from one end of the range to the other.

    `layoutA` and `layoutB` are required and should generally be set to
    `#nvvm.mma_layout<row>` and `#nvvm.mma_layout<col>` respectively, but other
    combinations are possible for certain layouts according to the table below.

    ```
    | A/B Type | Shape     | ALayout | BLayout | A Type   | B Type   | C/D Type          |
    |----------|-----------|---------|---------|----------|----------|-------------------|
    | f64      | .m8n8k4   | row     | col     | 1x f64   | 1x f64   | 2x f64            |
    | f16      | .m8n8k4   | row/col | row/col | 2x f16x2 | 2x f16x2 | 4x f16x2 or 8xf32 |
    |          | .m16n8k8  | row     | col     | 2x f16x2 | 1x f16x2 | 2x f16x2 or 4 f32 |
    |          | .m16n8k16 | row     | col     | 4x f16x2 | 2x f16x2 | 2x f16x2 or 4 f32 |
    | bf16     | .m16n8k8  | row     | col     | 2x i32   | 1x i32   | 4x f32            |
    |          | .m16n8k16 | row     | col     | 4x i32   | 2x i32   | 4x f32            |
    | tf32     | .m16n8k4  | row     | col     | 2x i32   | 1x i32   | 4x f32            |
    |          | .m16n8k8  | row     | col     | 4x i32   | 2x i32   | 2x f16x2 or 4 f32 |
    | u8/s8    | .m8n8k16  | row     | col     | 1x i32   | 1x i32   | 2x i32            |
    |          | .m16n8k16 | row     | col     | 2x i32   | 1x i32   | 4x i32            |
    |          | .m16n8k32 | row     | col     | 4x i32   | 2x i32   | 4x i32            |
    | u4/s4    | .m8n8k32  | row     | col     | 1x i32   | 1x i32   | 2x i32            |
    |          | m16n8k32  | row     | col     | 2x i32   | 1x i32   | 4x i32            |
    |          | m16n8k64  | row     | col     | 4x i32   | 2x i32   | 4x i32            |
    | b1       | m8n8k128  | row     | col     | 1x i32   | 1x i32   | 2x i32            |
    |          | m16n8k128 | row     | col     | 2x i32   | 1x i32   | 4x i32            |
    ```


    Example:
    ```mlir

    %128 = nvvm.mma.sync A[%120, %121, %122, %123]
                         B[%124, %125]
                         C[%126, %127]
                         {layoutA = #nvvm.mma_layout<row>,
                          layoutB = #nvvm.mma_layout<col>,
                          shape = {k = 16 : i32, m = 16 : i32, n = 8 : i32}}
        : (vector<2xf16>, vector<2xf16>, vector<2xf16>)
           -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
    ```
  }];

  let results = (outs LLVM_AnyStruct:$res);
  let arguments = (ins NVVM_MMAShapeAttr:$shape,
             OptionalAttr<MMAB1OpAttr>:$b1Op,
             OptionalAttr<MMAIntOverflowAttr>:$intOverflowBehavior,
             MMALayoutAttr:$layoutA,
             MMALayoutAttr:$layoutB,
             OptionalAttr<MMATypesAttr>:$multiplicandAPtxType,
             OptionalAttr<MMATypesAttr>:$multiplicandBPtxType,
             Variadic<LLVM_Type>:$operandA,
             Variadic<LLVM_Type>:$operandB,
             Variadic<LLVM_Type>:$operandC);

  let extraClassDeclaration = !strconcat([{
      static llvm::Intrinsic::ID getIntrinsicID(
            int64_t m, int64_t n, uint64_t k,
            std::optional<MMAB1Op> b1Op,
            std::optional<MMAIntOverflow> sat,
            mlir::NVVM::MMALayout layoutAEnum, mlir::NVVM::MMALayout layoutBEnum,
            mlir::NVVM::MMATypes eltypeAEnum, mlir::NVVM::MMATypes eltypeBEnum,
            mlir::NVVM::MMATypes eltypeCEnum, mlir::NVVM::MMATypes eltypeDEnum) {
        llvm::StringRef layoutA = stringifyEnum(layoutAEnum);
        llvm::StringRef layoutB = stringifyEnum(layoutBEnum);
        llvm::StringRef eltypeA = stringifyEnum(eltypeAEnum);
        llvm::StringRef eltypeB = stringifyEnum(eltypeBEnum);
        llvm::StringRef eltypeC = stringifyEnum(eltypeCEnum);
        llvm::StringRef eltypeD = stringifyEnum(eltypeDEnum);
        }],
        MMA_SYNC_INTR<>.id, [{
        return 0;
      }

      static std::optional<mlir::NVVM::MMATypes> inferOperandMMAType(Type operandElType,
        bool isAccumulator);

      MMATypes accumPtxType();
      MMATypes resultPtxType();
    }]);

  let builders = [
      OpBuilder<(ins  "Type":$resultType, "ValueRange":$operandA,
        "ValueRange":$operandB, "ValueRange":$operandC,
        "ArrayRef<int64_t>":$shape, "std::optional<MMAB1Op>":$b1Op,
        "std::optional<MMAIntOverflow>":$intOverflow,
        "std::optional<std::array<MMATypes, 2>>":$multiplicandPtxTypes,
        "std::optional<std::array<MMALayout, 2>>":$multiplicandLayouts)>
    ];

  string llvmBuilder = [{
    auto operands = moduleTranslation.lookupValues(opInst.getOperands());
    auto intId = mlir::NVVM::MmaOp::getIntrinsicID(
        $shape.getM(), $shape.getN(), $shape.getK(),
        $b1Op, $intOverflowBehavior,
        $layoutA, $layoutB,
        *$multiplicandAPtxType,
        *$multiplicandBPtxType,
        op.accumPtxType(),
        op.resultPtxType());

    $res = createIntrinsicCall(
      builder, intId, operands);
  }];

  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;
}

/// Generate enum value of the mma.sync intrinsic.
class MMA_SP_SYNC_NAME<string Metadata, string Kind, int Satfinite,
                       WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  string signature = MMA_SIGNATURE<A, B, C, D>.ret;
  string id = "llvm::Intrinsic::nvvm_mma"
              # "_" # !subst("::", "_", Metadata)
              # "_" # A.geom
              # "_row_col"
              # !if(!ne(Kind, ""), !strconcat("_", !subst("::", "_", Kind)), "")
              # !if(Satfinite, "_satfinite", "")
              # signature;
}

// Returns true if this combination of layout/kind/satf for MMA.SP ops is supported;
// false otherwise.
// E.g.
// if NVVM_MMA_SP_SUPPORTED<...>.ret then
//   def : FOO<>; // The record will only be defined for supported ops.
//
class NVVM_MMA_SP_SUPPORTED<list<WMMA_REGS> frags, string metadata,
                            string kind, int satf> {
  // MMA.SP ops check both layouts.
  string a_type = frags[0].ptx_elt_type;
  string b_type = frags[1].ptx_elt_type;
  string c_type = frags[2].ptx_elt_type;
  string d_type = frags[3].ptx_elt_type;
  string geom = frags[0].geom;

  bit is_int = !or(!eq(a_type, "s8"),
                   !eq(a_type, "u8"),
                   !eq(a_type, "s4"),
                   !eq(a_type, "u4"));

  bit ret = !cond(

    // Limit satf to valid types
    !and(!eq(satf, 1),
         !eq(is_int, 0)): false,

    // f16/bf16/tf32 requires A and B to be the same type.
    !and(!or(!eq(a_type, "f16"),
             !eq(a_type, "bf16"),
             !eq(a_type, "tf32")),
         !ne(a_type, b_type)): false,

    // m16n8k16, m16n8k32 and m16n8k64 requires C and D to be the same type.
    !and(!or(!eq(geom, "m16n8k16"),
             !eq(geom, "m16n8k32"),
             !eq(geom, "m16n8k64")),
         !ne(c_type, d_type)): false,

    !and(!eq(kind, ""),
         !or(!eq(a_type, "e3m2"),
             !eq(a_type, "e2m3"),
             !eq(a_type, "e2m1"),
             !eq(b_type, "e3m2"),
             !eq(b_type, "e2m3"),
             !eq(b_type, "e2m1"))): false,

    !and(!eq(kind, ""),
         !eq(geom, "m16n8k64"),
         !or(!eq(c_type, "f16"),
             !eq(d_type, "f16"))): false,

    !and(!ne(kind, ""),
         !or(!eq(metadata, "sp"),
             !ne(geom, "m16n8k64"),
             !eq(is_int, 1))): false,

    // All other are OK.
    true: true
  );
}

/// Helper to create the mapping between the configuration and the mma.sp.sync
/// intrinsic enum value.
class MMA_SP_SYNC_INTR {
  list<list<list<list<string>>>> cond0 =
    !foreach(op, NVVM_MMA_OPS.all_mma_sp_sync_ops,
      !foreach(metadata, ["sp", "sp::ordered_metadata"],
        !foreach(kind, ["", "kind::f8f6f4"],
          !foreach (satf, [0, 1],
            !if(NVVM_MMA_SP_SUPPORTED<op, metadata, kind, satf>.ret,
                "if (m == " # op[0].m # " && n == " # op[0].n # " && k == " # op[0].k
                # " && \"" # op[0].ptx_elt_type # "\" == eltypeA"
                # " && \"" # op[1].ptx_elt_type # "\" == eltypeB"
                # " && \"" # op[2].ptx_elt_type # "\" == eltypeC"
                # " && \"" # op[3].ptx_elt_type # "\" == eltypeD"
                # " && (satf.has_value()  ? " # satf # " == static_cast<int>(*satf) : true)"
                # " && " # !if(!eq(metadata, "sp"), "!orderedMetadata", "orderedMetadata") # ")\n"
                # "  return " #
                MMA_SP_SYNC_NAME<metadata, kind, satf, op[0], op[1], op[2], op[3]>.id # ";",
                "") // if supported
          ) // satf
        ) // kind
      ) // metadata
    ); // all_mma_sp_sync_ops
  list<list<list<string>>> f1 = !foldl([[[""]]], cond0, acc, el,
                                       !listconcat(acc, el));
  list<list<string>> f2 = !foldl([[""]], f1, acc, el, !listconcat(acc, el));
  list<string> f3 = !foldl([""], f2, acc, el, !listconcat(acc, el));
  string id = !foldl("", f3, acc, el, acc # "\n" # el);
}

def NVVM_MmaSpOp : NVVM_Op<"mma.sp.sync", [AttrSizedOperandSegments]> {

  let summary = "cooperative sparse matrix-multiply and accumulate";

  let description = [{
    The `nvvm.mma.sp.sync` operation collectively performs the sparse operation
    `D = matmul(A_sparse, B) + C` using all threads in a warp.

    This operation is similar to `nvvm.mma.sync` but with structured sparsity
    in the A operand. The sparsity follows the 2:4 structured sparse pattern
    where 2 out of every 4 elements are non-zero.

    All the threads in the warp must execute the same `mma.sp.sync` operation.

    The `sparseMetadata` operand provides the sparsity indices that indicate
    which elements in the A operand are non-zero. The `sparsitySelector`
    controls how the indices are distributed among threads in the warp and
    should typically be 0 or 1.

    The optional `orderedMetadata` attribute specifies the metadata ordering:
    - Absence (default): Uses standard sparse metadata ordering
    - Presence: Uses ordered metadata (PTX ISA 8.5+, sm_90+)

    The optional `kind` attribute specifies mixed-precision modes for FP8 operations:
    - `f8f6f4`: Enables e3m2, e2m3, e2m1 FP8 types and f16 accumulator (PTX ISA 8.7+, sm_90+)
    - Only valid with ordered metadata and m16n8k64 shape

    The shapes, layouts, and data types follow the same constraints as the
    regular `nvvm.mma.sync` operation, but the A operand contains only the
    non-zero elements in compressed format.

    Example:
    ```mlir
    %d = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]
                          sparseMetadata[%meta] selector[%sel]
                          {shape = {k = 32 : i32, m = 16 : i32, n = 8 : i32}}
        : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>

    // With ordered metadata:
    %d = nvvm.mma.sp.sync A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]
                          sparseMetadata[%meta] selector[%sel]
                          {orderedMetadata, shape = {k = 32 : i32, m = 16 : i32, n = 8 : i32}}
        : (vector<2xf16>, vector<2xf16>, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
    ```
  }];

  let results = (outs LLVM_AnyStruct:$res);
  let arguments = (ins NVVM_MMAShapeAttr:$shape,
             OptionalAttr<MMAIntOverflowAttr>:$intOverflowBehavior,
             OptionalAttr<MMATypesAttr>:$multiplicandAPtxType,
             OptionalAttr<MMATypesAttr>:$multiplicandBPtxType,
             UnitAttr:$orderedMetadata,
             OptionalAttr<MMAKindAttr>:$kind,
             Variadic<LLVM_Type>:$operandA,
             Variadic<LLVM_Type>:$operandB,
             Variadic<LLVM_Type>:$operandC,
             I32:$sparseMetadata,
             I32:$sparsitySelector);

  let extraClassDeclaration = !strconcat([{
      static llvm::Intrinsic::ID getIntrinsicID(
            int64_t m, int64_t n, uint64_t k,
            std::optional<MMAIntOverflow> satf,
            bool orderedMetadata,
            std::optional<MMAKind> kind,
            mlir::NVVM::MMATypes eltypeAEnum, mlir::NVVM::MMATypes eltypeBEnum,
            mlir::NVVM::MMATypes eltypeCEnum, mlir::NVVM::MMATypes eltypeDEnum) {
        llvm::StringRef eltypeA = stringifyEnum(eltypeAEnum);
        llvm::StringRef eltypeB = stringifyEnum(eltypeBEnum);
        llvm::StringRef eltypeC = stringifyEnum(eltypeCEnum);
        llvm::StringRef eltypeD = stringifyEnum(eltypeDEnum);
        }],
        MMA_SP_SYNC_INTR<>.id, [{
          return 0;
      }

      static std::optional<mlir::NVVM::MMATypes> inferOperandMMAType(Type operandElType,
        bool isAccumulator);

      MMATypes accumPtxType();
      MMATypes resultPtxType();

      static mlir::NVVM::IDArgPair
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::IRBuilderBase& builder);
    }]);

  let builders = [
      OpBuilder<(ins  "Type":$resultType, "ValueRange":$operandA,
        "ValueRange":$operandB, "ValueRange":$operandC,
        "Value":$sparseMetadata, "Value":$sparsitySelector,
        "ArrayRef<int64_t>":$shape,
        "std::optional<MMAIntOverflow>":$intOverflow,
        "std::optional<std::array<MMATypes, 2>>":$multiplicandPtxTypes)>
    ];

  string llvmBuilder = [{
    auto [id, args] = NVVM::MmaSpOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    $res = createIntrinsicCall(builder, id, args);
  }];

  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;
}

def ScaleVecSize1X  : I32EnumAttrCase<"X1", 0, "x1">;
def ScaleVecSize2X  : I32EnumAttrCase<"X2", 1, "x2">;
def ScaleVecSize4X  : I32EnumAttrCase<"X4", 2, "x4">;

def ScaleVecSize : I32EnumAttr<
  "ScaleVecSize",
  "MMA Scale Vector Sizes",
  [ScaleVecSize1X, ScaleVecSize2X, ScaleVecSize4X]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}

def ScaleVecSizeAttr : EnumAttr<NVVM_Dialect, ScaleVecSize, "scale_vec_size"> {
  let assemblyFormat = "`<` $value `>`";
}

def UE8M0 : I32EnumAttrCase<"UE8M0", 0, "ue8m0">;
def UE4M3 : I32EnumAttrCase<"UE4M3", 1, "ue4m3">;

def BlockScaleFormat : I32EnumAttr<
  "BlockScaleFormat",
  "MMA Block Scale Format",
  [UE8M0, UE4M3]
> {
  let cppNamespace = "::mlir::NVVM";
  let genSpecializedAttr = 0;
}

def BlockScaleFormatAttr : EnumAttr<NVVM_Dialect, BlockScaleFormat, "block_scale_format"> {
  let assemblyFormat = "`<` $value `>`";
}

def MMABlockScaleKindMXF8F6F4  : I32EnumAttrCase<"MXF8F6F4", 0, "mxf8f6f4">;
def MMABlockScaleKindMXF4  : I32EnumAttrCase<"MXF4", 1, "mxf4">;
def MMABlockScaleKindMXF4NVF4  : I32EnumAttrCase<"MXF4NVF4", 2, "mxf4nvf4">;

def MMABlockScaleKind : I32EnumAttr<
  "MMABlockScaleKind",
  "Block Scale Kind",
  [MMABlockScaleKindMXF8F6F4, MMABlockScaleKindMXF4, MMABlockScaleKindMXF4NVF4]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}

def MMABlockScaleKindAttr : EnumAttr<NVVM_Dialect, MMABlockScaleKind, "block_scale_kind"> {
  let description = [{
    The MMABlockScaleKind attribute describes the allowed set of types for matrix A and B in the *.mma.{sp}.block_scale Op. The following are supported types for each kind:

    ```
    +--------------+-------------------------------------------+
    | Matrix Kind  |      supported types for A / B            |
    +--------------+-------------------------------------------+
    | mxf8f6f4     | e4m3, e5m2, e2m3, e3m2, e2m1              |
    | mxf4         | e2m1                                      |
    | mxf4nvf4     | e2m1                                      |
    +--------------+-------------------------------------------+
    ```
  }];
  let assemblyFormat = "`<` $value `>`";
}

/// Generate enum value of the mma.block_scale intrinsic.
class MMA_BLOCK_SCALE_NAME<string Kind, string SType, string ScaleVecSize,
                           WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  string signature = MMA_SIGNATURE<A, B, C, D>.ret;
  string id = "llvm::Intrinsic::nvvm_mma_block_scale"
              # "_" # A.geom
              # "_row_col"
              # "_" # Kind
              # !subst(".", "_", ScaleVecSize)
              # signature
              # "_" # SType;
}

/// Generate enum value of the mma.sp.block_scale intrinsic.
class MMA_SP_BLOCK_SCALE_NAME<string Kind, string SType, string ScaleVecSize,
                              WMMA_REGS A, WMMA_REGS B, WMMA_REGS C, WMMA_REGS D> {
  string signature = MMA_SIGNATURE<A, B, C, D>.ret;
  string id = "llvm::Intrinsic::nvvm_mma_sp_ordered_metadata_block_scale"
              # "_" # A.geom
              # "_row_col"
              # "_" # Kind
              # !subst(".", "_", ScaleVecSize)
              # signature
              # "_" # SType;
}

// Returns true if this combination is supported for MMA.BLOCK_SCALE ops.
// This references the NVVM_MMA_BLOCK_SCALE_SUPPORTED class from IntrinsicsNVVM.td
class NVVM_MMA_BLOCK_SCALE_SUPPORTED<list<WMMA_REGS> frags, string kind,
                                     string stype, string scale_vec_size> {
  string geom = frags[0].geom;
  bit ret = !cond(
    !and(!eq(geom, "m16n8k64"),
         !eq(kind, "mxf4"),
         !or(!eq(scale_vec_size, ""),
             !eq(scale_vec_size, ".scale_2x")),
         !eq(stype, "ue8m0")) : true,
    !and(!eq(geom, "m16n8k64"),
         !eq(kind, "mxf4nvf4"),
         !eq(scale_vec_size, ".scale_2x"),
         !eq(stype, "ue8m0")) : true,
    !and(!eq(geom, "m16n8k64"),
         !eq(kind, "mxf4nvf4"),
         !eq(scale_vec_size, ".scale_4x"),
         !or(!eq(stype, "ue4m3"),
             !eq(stype, "ue8m0"))) : true,
    !and(!eq(geom, "m16n8k32"),
         !eq(kind, "mxf8f6f4"),
         !or(!eq(scale_vec_size, ""),
             !eq(scale_vec_size, ".scale_1x")),
         !eq(stype, "ue8m0")) : true,
    true: false
  );
}

// Returns true if this combination is supported for MMA.SP.BLOCK_SCALE ops.
// This references the NVVM_MMA_SP_BLOCK_SCALE_SUPPORTED class from IntrinsicsNVVM.td
class NVVM_MMA_SP_BLOCK_SCALE_SUPPORTED<list<WMMA_REGS> frags, string kind,
                                        string stype, string scale_vec_size> {
  string geom = frags[0].geom;
  bit ret = !cond(
    !and(!eq(geom, "m16n8k128"),
         !eq(kind, "mxf4"),
         !eq(stype, "ue8m0"),
         !or(!eq(scale_vec_size, ""),
             !eq(scale_vec_size, ".scale_2x"))): true,
    !and(!eq(geom, "m16n8k128"),
         !eq(kind, "mxf4nvf4"),
         !eq(stype, "ue8m0"),
         !eq(scale_vec_size, ".scale_2x")): true,
    !and(!eq(geom, "m16n8k128"),
         !eq(kind, "mxf4nvf4"),
         !or(!eq(stype, "ue4m3"),
             !eq(stype, "ue8m0")),
         !eq(scale_vec_size, ".scale_4x")): true,
    !and(!eq(geom, "m16n8k64"),
         !eq(kind, "mxf8f6f4"),
         !eq(stype, "ue8m0"),
         !or(!eq(scale_vec_size, ""),
             !eq(scale_vec_size, ".scale_1x"))): true,
    true: false
  );
}

/// Helper to create the mapping between the configuration and the mma.block_scale
/// intrinsic enum value.
class MMA_BLOCK_SCALE_INTR {
  list<list<list<list<string>>>> cond0 =
    !foreach(op, NVVM_MMA_OPS.all_mma_block_scale_ops,
      !foreach(kind, ["mxf4", "mxf4nvf4", "mxf8f6f4"],
        !foreach(scale_vec_size, ["", ".scale_1x", ".scale_2x", ".scale_4x"],
          !foreach(stype, ["ue8m0", "ue4m3"],
            !if(NVVM_MMA_BLOCK_SCALE_SUPPORTED<op, kind, stype, scale_vec_size>.ret,
                "if (m == " # op[0].m # " && n == " # op[0].n # " && k == " # op[0].k
                # " && \"" # op[0].ptx_elt_type # "\" == eltypeA"
                # " && \"" # op[1].ptx_elt_type # "\" == eltypeB"
                # " && \"" # op[2].ptx_elt_type # "\" == eltypeC"
                # " && \"" # kind # "\" == stringifyEnum(kind)"
                # " && \"" # stype # "\" == stringifyEnum(blockScaleFormat)"
                # " && \"" # scale_vec_size # "\" == getScaleVecSizeStr(scaleVecSize))\n"
                # "  return " #
                MMA_BLOCK_SCALE_NAME<kind, stype, scale_vec_size, op[0], op[1], op[2], op[3]>.id # ";",
                "") // if supported
          ) // stype
        ) // scale_vec_size
      ) // kind
    ); // all_mma_block_scale_ops
  list<list<list<string>>> f1 = !foldl([[[""]]], cond0, acc, el,
                                       !listconcat(acc, el));
  list<list<string>> f2 = !foldl([[""]], f1, acc, el, !listconcat(acc, el));
  list<string> f3 = !foldl([""], f2, acc, el, !listconcat(acc, el));
  string id = !foldl("", f3, acc, el, acc # "\n" # el);
}

/// Helper to create the mapping between the configuration and the mma.sp.block_scale
/// intrinsic enum value.
class MMA_SP_BLOCK_SCALE_INTR {
  list<list<list<list<string>>>> cond0 =
    !foreach(op, NVVM_MMA_OPS.all_mma_sp_block_scale_ops,
      !foreach(kind, ["mxf4", "mxf4nvf4", "mxf8f6f4"],
        !foreach(scale_vec_size, ["", ".scale_1x", ".scale_2x", ".scale_4x"],
          !foreach(stype, ["ue8m0", "ue4m3"],
            !if(NVVM_MMA_SP_BLOCK_SCALE_SUPPORTED<op, kind, stype, scale_vec_size>.ret,
                "if (m == " # op[0].m # " && n == " # op[0].n # " && k == " # op[0].k
                # " && \"" # op[0].ptx_elt_type # "\" == eltypeA"
                # " && \"" # op[1].ptx_elt_type # "\" == eltypeB"
                # " && \"" # op[2].ptx_elt_type # "\" == eltypeC"
                # " && \"" # kind # "\" == stringifyEnum(kind)"
                # " && \"" # stype # "\" == stringifyEnum(blockScaleFormat)"
                # " && \"" # scale_vec_size # "\" == getScaleVecSizeStr(scaleVecSize))\n"
                # "  return " #
                MMA_SP_BLOCK_SCALE_NAME<kind, stype, scale_vec_size, op[0], op[1], op[2], op[3]>.id # ";",
                "") // if supported
          ) // stype
        ) // scale_vec_size
      ) // kind
    ); // all_mma_sp_block_scale_ops
  list<list<list<string>>> f1 = !foldl([[[""]]], cond0, acc, el,
                                       !listconcat(acc, el));
  list<list<string>> f2 = !foldl([[""]], f1, acc, el, !listconcat(acc, el));
  list<string> f3 = !foldl([""], f2, acc, el, !listconcat(acc, el));
  string id = !foldl("", f3, acc, el, acc # "\n" # el);
}

// Common base class for MMA block scale operations (dense and sparse)
class NVVM_MmaBlockScaleBase<string mnemonic, list<Trait> traits = []> :
    NVVM_Op<mnemonic, !listconcat([AttrSizedOperandSegments], traits)> {

  let results = (outs LLVM_AnyStruct:$res);

  // Common attributes shared by both dense and sparse variants
  dag commonArguments = (ins
           NVVM_MMAShapeAttr:$shape,
           OptionalAttr<MMATypesAttr>:$multiplicandAPtxType,
           OptionalAttr<MMATypesAttr>:$multiplicandBPtxType,
           ScaleVecSizeAttr:$scaleVecSize,
           BlockScaleFormatAttr:$blockScaleFormat,
           MMABlockScaleKindAttr:$kind);

  // Common variadic operands for A, B, C matrices
  dag commonVariadicOperands = (ins
           Variadic<LLVM_Type>:$operandA,
           Variadic<LLVM_Type>:$operandB,
           Variadic<LLVM_Type>:$operandC);

  // Common scale operands for both A and B
  dag commonScaleOperands = (ins
             I32:$scaleAData,
             I16:$byteIdA,
             I16:$threadIdA,
             I32:$scaleBData,
             I16:$byteIdB,
             I16:$threadIdB);

  let extraClassDeclaration = !strconcat([{
      static llvm::Intrinsic::ID getIntrinsicID(
            int64_t m, int64_t n, uint64_t k,
            mlir::NVVM::MMATypes eltypeAEnum, mlir::NVVM::MMATypes eltypeBEnum,
            mlir::NVVM::MMATypes eltypeCEnum,
            mlir::NVVM::ScaleVecSize scaleVecSize,
            mlir::NVVM::BlockScaleFormat blockScaleFormat,
            mlir::NVVM::MMABlockScaleKind kind) {
        llvm::StringRef eltypeA = stringifyEnum(eltypeAEnum);
        llvm::StringRef eltypeB = stringifyEnum(eltypeBEnum);
        llvm::StringRef eltypeC = stringifyEnum(eltypeCEnum);

        auto getScaleVecSizeStr = [](ScaleVecSize svs) -> std::string {
          switch (svs) {
            case ScaleVecSize::X1: return ".scale_1x";
            case ScaleVecSize::X2: return ".scale_2x";
            case ScaleVecSize::X4: return ".scale_4x";
          }
          return "";
        };
        }],
        MMA_BLOCK_SCALE_INTR<>.id, [{
        return 0;
      }

      static mlir::NVVM::IDArgPair
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::IRBuilderBase& builder);
    }]);

  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;
}

def NVVM_MmaBlockScaleOp : NVVM_MmaBlockScaleBase<"mma.block_scale"> {

  let summary = "cooperative matrix-multiply and accumulate with block scaling";

  let description = [{
    The `nvvm.mma.block_scale` operation collectively performs the operation
    `D = matmul(A * SF_A, B * SF_B) + C` using all threads in a warp.

    A, B, C and D are dense matrices and SF_A and SF_B are scaling factors.
    Dimensions of SF_A and SF_B are based on scale vector sizes (x1, x2, x4),
    and the data type must be either ue8m0 or ue4m3.

    All the threads in the warp must execute the same `mma.block_scale` operation.

    This operation follows the same design pattern as `nvvm.mma.sync`, with additional
    scaling operands for both A and B matrices.

    Example:
    ```mlir
    %d = nvvm.mma.block_scale A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]
                              scaleA[%scaleAData, %byteIdA, %threadIdA]
                              scaleB[%scaleBData, %byteIdB, %threadIdB]
                              {shape = #nvvm.shape<m = 16, n = 8, k = 64>,
                               multiplicandAPtxType = #nvvm.mma_type<e2m1>,
                               multiplicandBPtxType = #nvvm.mma_type<e2m1>,
                               scaleVecSize = #nvvm.scale_vec_size<x2>,
                               blockScaleFormat = #nvvm.block_scale_format<ue8m0>,
                               kind = #nvvm.block_scale_kind<mxf4nvf4>}
        : (vector<4xf16>, vector<2xf16>, vector<2xf32>) -> !llvm.struct<(f32, f32)>
    ```
  }];

  // Combine common attributes and operands
  let arguments = !con(commonArguments, commonVariadicOperands, commonScaleOperands);

  let builders = [
      OpBuilder<(ins "Type":$resultType, "ValueRange":$operandA,
        "ValueRange":$operandB, "ValueRange":$operandC,
        "Value":$scaleAData, "Value":$byteIdA, "Value":$threadIdA,
        "Value":$scaleBData, "Value":$byteIdB, "Value":$threadIdB,
        "ArrayRef<int64_t>":$shape,
        "std::optional<std::array<MMATypes, 2>>":$multiplicandPtxTypes,
        "ScaleVecSize":$scaleVecSize,
        "BlockScaleFormat":$blockScaleFormat,
        "MMABlockScaleKind":$kind)>
    ];

  string llvmBuilder = [{
    auto [id, args] = NVVM::MmaBlockScaleOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    $res = createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_MmaSpBlockScaleOp : NVVM_MmaBlockScaleBase<"mma.sp.block_scale"> {

  let summary = "cooperative sparse matrix-multiply and accumulate with block scaling";

  let description = [{
    The `nvvm.mma.sp.block_scale` operation collectively performs the operation
    `D = matmul(A_sparse * SF_A, B * SF_B) + C` using all threads in a warp.

    A is a sparse matrix, and B, C and D are dense matrices.
    SF_A and SF_B are scaling factors.
    Dimensions of SF_A and SF_B are based on scale vector sizes (x1, x2, x4),
    and the data type must be either ue8m0 or ue4m3.

    This operation is similar to `nvvm.mma.block_scale` but with structured sparsity
    in the A operand. The sparsity follows the 2:4 structured sparse pattern
    where 2 out of every 4 elements are non-zero.

    All the threads in the warp must execute the same `mma.sp.block_scale` operation.

    The `sparseMetadata` operand provides the sparsity indices that indicate
    which elements in the A operand are non-zero. The `sparsitySelector`
    controls how the indices are distributed among threads in the warp and
    should typically be 0 or 1.

    This operation follows the same design pattern as `nvvm.mma.sp.sync`, with additional
    scaling operands for both A and B matrices. Note that sparse block scale operations
    always use ordered metadata (sm_90+).

    Example:
    ```mlir
    %d = nvvm.mma.sp.block_scale A[%a0, %a1] B[%b0, %b1] C[%c0, %c1]
                                 sparseMetadata[%meta] selector[%sel]
                                 scaleA[%scaleAData, %byteIdA, %threadIdA]
                                 scaleB[%scaleBData, %byteIdB, %threadIdB]
                                 {shape = #nvvm.shape<m = 16, n = 8, k = 128>,
                                  multiplicandAPtxType = #nvvm.mma_type<e2m1>,
                                  multiplicandBPtxType = #nvvm.mma_type<e2m1>,
                                  scaleVecSize = #nvvm.scale_vec_size<x2>,
                                  blockScaleFormat = #nvvm.block_scale_format<ue8m0>,
                                  kind = #nvvm.block_scale_kind<mxf4>}
        : (vector<2xf16>, vector<2xf16>, vector<2xf32>) -> !llvm.struct<(f32, f32)>
    ```
  }];

  // Sparse-specific attributes and operands
  dag sparseSpecificArguments = (ins
           UnitAttr:$orderedMetadata);

  dag sparseSpecificOperands = (ins
             I32:$sparseMetadata,
             I32:$sparsitySelector);

  // Combine common and sparse-specific attributes and operands
  let arguments = !con(commonArguments, sparseSpecificArguments, 
                       commonVariadicOperands, sparseSpecificOperands,
                       commonScaleOperands);

  // Override extraClassDeclaration to use sparse intrinsics
  let extraClassDeclaration = !strconcat([{
      static llvm::Intrinsic::ID getIntrinsicID(
            int64_t m, int64_t n, uint64_t k,
            mlir::NVVM::MMATypes eltypeAEnum, mlir::NVVM::MMATypes eltypeBEnum,
            mlir::NVVM::MMATypes eltypeCEnum,
            mlir::NVVM::ScaleVecSize scaleVecSize,
            mlir::NVVM::BlockScaleFormat blockScaleFormat,
            mlir::NVVM::MMABlockScaleKind kind) {
        llvm::StringRef eltypeA = stringifyEnum(eltypeAEnum);
        llvm::StringRef eltypeB = stringifyEnum(eltypeBEnum);
        llvm::StringRef eltypeC = stringifyEnum(eltypeCEnum);

        auto getScaleVecSizeStr = [](ScaleVecSize svs) -> std::string {
          switch (svs) {
            case ScaleVecSize::X1: return ".scale_1x";
            case ScaleVecSize::X2: return ".scale_2x";
            case ScaleVecSize::X4: return ".scale_4x";
          }
          return "";
        };
        }],
        MMA_SP_BLOCK_SCALE_INTR<>.id, [{
        return 0;
      }

      static mlir::NVVM::IDArgPair
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::IRBuilderBase& builder);
    }]);

  let builders = [
      OpBuilder<(ins "Type":$resultType, "ValueRange":$operandA,
        "ValueRange":$operandB, "ValueRange":$operandC,
        "Value":$sparseMetadata, "Value":$sparsitySelector,
        "Value":$scaleAData, "Value":$byteIdA, "Value":$threadIdA,
        "Value":$scaleBData, "Value":$byteIdB, "Value":$threadIdB,
        "ArrayRef<int64_t>":$shape,
        "std::optional<std::array<MMATypes, 2>>":$multiplicandPtxTypes,
        "ScaleVecSize":$scaleVecSize,
        "BlockScaleFormat":$blockScaleFormat,
        "MMABlockScaleKind":$kind)>
    ];

  string llvmBuilder = [{
    auto [id, args] = NVVM::MmaSpBlockScaleOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    $res = createIntrinsicCall(builder, id, args);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM TMA Ops
//===----------------------------------------------------------------------===//

// List of modes supported for TMA Load and Prefetch Ops
def TMALoadModeTile   : I32EnumAttrCase<"TILE", 0, "tile">;
def TMALoadModeIm2Col : I32EnumAttrCase<"IM2COL", 1, "im2col">;
def TMALoadModeIm2ColW : I32EnumAttrCase<"IM2COL_W", 2, "im2col_w">;
def TMALoadModeIm2ColW128 : I32EnumAttrCase<"IM2COL_W_128", 3, "im2col_w_128">;
def TMALoadModeTileGather4 : I32EnumAttrCase<"TILE_GATHER4", 4, "tile_gather4">;

def TMALoadMode : I32EnumAttr<"TMALoadMode", "NVVM TMA Load Mode",
    [TMALoadModeTile, TMALoadModeIm2Col,
     TMALoadModeIm2ColW, TMALoadModeIm2ColW128,
     TMALoadModeTileGather4]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def TMALoadModeAttr : EnumAttr<NVVM_Dialect, TMALoadMode, "tma_load_mode"> {
  let summary = "List of Load-Modes supported for TMA Tensor Ops";
  let description = [{
    TMA Tensor Ops support the following modes, when copying data from
    global memory to shared memory (i.e. load):

    Tile Mode: It's the default mode. The source multi-dimensional tensor
    layout is preserved at the destination.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-tiled-mode)

    Im2col Mode: This mode is used when `im2colOffsets` operands are present.
    The elements in the Bounding Box of the source tensor are rearranged into
    columns at the destination. In this mode, the tensor has to be at least
    3-dimensional. The number of `im2colOffsets` is `dims - 2` where `dims`
    is the dimension of the tensor.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-im2col-mode)

    Im2col_W Mode: This mode is similar to Im2Col mode with the restriction that
    elements are accessed across the W dimension only. The number of `im2colOffsets`
    are always two, referred as `wHalo` and `wOffset`.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-im2col-w-w128-modes)

    Im2col_W_128 Mode: This mode is similar to Im2Col_W mode with the number of
    elements accessed across the W dimension is always 128 only.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-im2col-w-w128-modes)

    Tile_Gather4 Mode: This mode is similar to Tile mode but works only on 2D tensor.
    In gather4 mode, four rows in the source 2D tensor are combined to form a single
    2D tensor at the destination. This mode requires five co-ordinates. The first one
    represents the column-index followed by four row indices.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-tiled-scatter4-gather4-modes)
  }];

  let assemblyFormat = "`<` $value `>`";
}

// List of modes supported for TMA Store and Reduction Ops
def TMAStoreModeTile   : I32EnumAttrCase<"TILE", 0, "tile">;
def TMAStoreModeIm2Col : I32EnumAttrCase<"IM2COL", 1, "im2col">;
def TMAStoreModeTileScatter4 : I32EnumAttrCase<"TILE_SCATTER4", 2, "tile_scatter4">;

def TMAStoreMode : I32EnumAttr<"TMAStoreMode", "NVVM TMA Store Mode",
    [TMAStoreModeTile, TMAStoreModeIm2Col, TMAStoreModeTileScatter4]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def TMAStoreModeAttr : EnumAttr<NVVM_Dialect, TMAStoreMode, "tma_store_mode"> {
  let assemblyFormat = "`<` $value `>`";
}

// Num CTAs in a group participating in the TMA/MMA operations.
// This corresponds to the "cta_group::1", "cta_group::2"
// modifiers in the PTX instructions.
def CTAGroup_1 : I32EnumAttrCase<"CTA_1", 0, "cta_1">;
def CTAGroup_2 : I32EnumAttrCase<"CTA_2", 1, "cta_2">;

def CTAGroupKind : I32EnumAttr<"CTAGroupKind",
                            "NVVM CTA group kind",
  [CTAGroup_1, CTAGroup_2]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def CTAGroupKindAttr :
  EnumAttr<NVVM_Dialect, CTAGroupKind, "cta_group"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_CpAsyncBulkCommitGroupOp : NVVM_Op<"cp.async.bulk.commit.group">,
  Arguments<(ins )> {
  let assemblyFormat = "attr-dict";
  let description = [{
    This Op commits all prior initiated but uncommitted cp.async.bulk
    instructions into a cp.async.bulk-group.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group)
  }];

  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_cp_async_bulk_commit_group);
  }];
}

def NVVM_CpAsyncBulkWaitGroupOp : NVVM_Op<"cp.async.bulk.wait_group", [NVVMRequiresSM<90>]>,
  Arguments<(ins 
    ConfinedAttr<I32Attr, [IntMinValue<0>]>:$group, 
    OptionalAttr<UnitAttr>:$read)> {
  let assemblyFormat = "$group attr-dict";
  let description = [{
    Op waits for completion of the most recent bulk async-groups.

    The `$group` operand tells waiting has to be done until for $group or fewer
    of the most recent bulk async-groups. If `$group` is 0, the op wait until 
    all the most recent bulk async-groups have completed.

    The `$read` indicates that the waiting has to be done until all the bulk 
    async operations in the specified bulk async-group have completed reading 
    from their source locations.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group)
  }];
  
  string llvmBuilder = [{
    auto intId = op.getRead() ?
      llvm::Intrinsic::nvvm_cp_async_bulk_wait_group_read :
      llvm::Intrinsic::nvvm_cp_async_bulk_wait_group;
    createIntrinsicCall(builder, intId, builder.getInt32($group));
  }];
}

def NVVM_CpAsyncBulkTensorGlobalToSharedClusterOp : 
  NVVM_Op<"cp.async.bulk.tensor.shared.cluster.global", 
  [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>, 
  AttrSizedOperandSegments, NVVMRequiresSM<90>]>,
  Arguments<(ins  AnyTypeOf<[LLVM_PointerShared, LLVM_PointerSharedCluster]>:$dstMem,
                  LLVM_PointerGeneric:$tmaDescriptor,
                  Variadic<I32>:$coordinates,
                  LLVM_PointerShared:$mbar,                  
                  Variadic<I16>:$im2colOffsets,
                  Optional<I16>:$multicastMask,
                  Optional<I64>:$l2CacheHint,
                  DefaultValuedAttr<TMALoadModeAttr, "TMALoadMode::TILE">:$mode,
                  DefaultValuedAttr<BoolAttr, "false">:$isCTAOnly,
                  OptionalAttr<CTAGroupKindAttr>:$group,
                  PtxPredicate:$predicate)> {
  let description = [{
    Initiates an asynchronous copy operation on the tensor data from global 
    memory to shared::cluster (or) shared::cta memory. This Op supports all
    the load modes specified in `TMALoadMode`.

    The `multicastMask` operand is optional. When it is present, the Op copies
    data from global memory to shared memory of multiple CTAs in the cluster.
    Operand `multicastMask` specifies the destination CTAs in the cluster such 
    that each bit position in the 16-bit `multicastMask` operand corresponds to
    the `nvvm.read.ptx.sreg.ctaid` of the destination CTA.     

    The `l2CacheHint` operand is optional, and it is used to specify cache 
    eviction policy that may be used during the memory access.
    
    When the `isCTAOnly` attribute is set to true, the destination is
    shared::cta only. Hence, `multicastMask` and `CTAGroup` are not applicable
    when `isCTAOnly` is true.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor)
  }];

  let assemblyFormat = [{ 
    $dstMem `,` 
    $tmaDescriptor `,` 
    $mbar `,` 
    `box` `[`$coordinates `]` 
    (`im2col` `[` $im2colOffsets^ `]` )?
    (`multicast_mask` `=` $multicastMask^ )?
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    (`predicate` `=` $predicate^)? 
    attr-dict  `:` type($dstMem) `,` type($tmaDescriptor)
  }];

  let extraClassDefinition = [{
    std::string $cppClass::getPtx() {
      int im2colDim = getIm2colOffsets().size();
      int dim = getCoordinates().size();
      std::string ptx = "cp.async.bulk.tensor.";
      ptx += std::to_string(dim) + "d.";
      ptx += "shared::cluster.global.mbarrier::complete_tx::bytes";      
      if(im2colDim) ptx += ".im2col";
      if(getMulticastMask()) ptx += ".multicast::cluster";      
      if(getL2CacheHint()) ptx += ".L2::cache_hint";
      
      auto preg = [](int r) { return "%" + std::to_string(r); };

      // Build Registers
      ptx += " [%0], [%1, {";
      int r = 2;      
      for(int i = 0; i < dim; i++) ptx += preg(r+i) + ",";
      ptx.pop_back(); r += dim;
      ptx += "} ], [%" + std::to_string(r++) + "]";
      if(im2colDim) {
        ptx += ",{";
        for(int i = 0; i < im2colDim; i++) ptx += preg(r+i) + ",";
        ptx.pop_back(); r += im2colDim;
        ptx += "}";
      }
      if(getMulticastMask()) ptx += ", " + preg(r++);
      if(getL2CacheHint()) ptx += ", " + preg(r++);
      ptx += ";";
      return ptx;
    }
  }];
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    bool hasIntrinsic() { return !getPredicate(); }

    bool getAsmValues(RewriterBase &rewriter,
      llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues);

    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::CpAsyncBulkTensorGlobalToSharedClusterOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_CpAsyncBulkTensorSharedCTAToGlobalOp : 
  NVVM_PTXBuilder_Op<"cp.async.bulk.tensor.global.shared.cta",
  [AttrSizedOperandSegments]>,
  Arguments<(ins LLVM_PointerGeneric:$tmaDescriptor,
                 LLVM_PointerShared:$srcMem,
                 Variadic<I32>:$coordinates,
                 Optional<I64>:$l2CacheHint,
                 DefaultValuedAttr<TMAStoreModeAttr, "TMAStoreMode::TILE">:$mode,
                 PtxPredicate:$predicate)> {
  let description = [{
    Initiates an asynchronous copy of the tensor data from shared::cta
    memory to global memory. This Op supports all the store modes specified in
    `TMAStoreMode`.

    The `l2CacheHint` operand is optional, and it is used to specify cache
    eviction policy that may be used during the memory access.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor)
  }];

  let assemblyFormat = [{ 
    $tmaDescriptor `,` 
    $srcMem `,` 
    `box` `[`$coordinates `]` 
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    (`,` `predicate` `=` $predicate^)?
    attr-dict `:` type($tmaDescriptor) `,` type($srcMem)
  }];

  let extraClassDeclaration = [{
    bool hasIntrinsic() { return !getPredicate(); }

    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];

  let extraClassDefinition = [{
    std::string $cppClass::getPtx() {
      int dim = getCoordinates().size();
      std::string ptx = "cp.async.bulk.tensor.";
      ptx += std::to_string(dim) + "d.";
      ptx += "global.shared::cta.bulk_group";
      if(dim == 1) ptx += " [%0, {%2} ], [%1];";
      if(dim == 2) ptx += " [%0, {%2, %3} ], [%1];";
      if(dim == 3) ptx += " [%0, {%2, %3, %4} ], [%1];";
      if(dim == 4) ptx += " [%0, {%2, %3, %4, %5} ], [%1];";
      if(dim == 5) ptx += " [%0, {%2, %3, %4, %5, %6} ], [%1];";
      return ptx;
    }
  }];
  let hasVerifier = 1;

  string llvmBuilder = [{
    auto [id, args] = NVVM::CpAsyncBulkTensorSharedCTAToGlobalOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM Prefetch Op
//===----------------------------------------------------------------------===//

def PrefetchCacheLevelL1 : I32EnumCase<"L1", 0, "L1">;
def PrefetchCacheLevelL2 : I32EnumCase<"L2", 1, "L2">;

def PrefetchCacheLevel : I32Enum<"PrefetchCacheLevel",
                                 "NVVM Prefetch Cache Level",
                                 [PrefetchCacheLevelL1, PrefetchCacheLevelL2]> {
  let cppNamespace = "::mlir::NVVM";
}

def PrefetchCacheLevelAttr : EnumAttr<NVVM_Dialect, PrefetchCacheLevel, "prefetch_cache_level"> {
  let assemblyFormat = "$value";
}

def NVVM_PrefetchOp : NVVM_Op<"prefetch",
    [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]> {
  let summary = "Brings the cache line containing an address into the specified cache level";
  let description = [{
    Prefetches the cache line containing the address given by `addr`. The 
    operand may be a global, local, or generic pointer. When `tensormap` is 
    specified, the operand may instead be a constant or generic pointer. If the 
    address maps to shared memory, the operation has no effect.

    At most one of `cacheLevel` or `tensormap` may be present. The `cacheLevel` 
    attribute selects the target cache level. When combined with `uniform`, the 
    prefetch is performed to the uniform cache, in which case `addr` must be a 
    generic pointer.

    When `tensormap` is used, the line containing `addr` is brought from the 
    constant or parameter state space for later use by `cp.async.bulk.tensor`. 
    If `in_param_space` is specified, the generic pointer is interpreted as 
    referring to the parameter state space.

    `uniform` can be specified after the `cacheLevel` to indicate that the 
    prefetch is performed to the specified uniform cache level. If `uniform` is 
    specified, `addr` must be a generic address pointer and no operation is 
    performed if `addr` maps to a `const`, `local`, or `shared` memory location.

    The `evictPriority` attribute is optional and specifies the cache eviction
    priority when `cacheLevel` is L2.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-prefetch-prefetchu)
  }];
  let arguments = (ins OptionalAttr<PrefetchCacheLevelAttr>:$cacheLevel,
                       OptionalAttr<CacheEvictionPriorityAttr>:$evictPriority,
                       AnyTypeOf<[LLVM_PointerGlobal,
                                  LLVM_PointerLocal,
                                  LLVM_PointerGeneric,
                                  LLVM_PointerConst]>:$addr,
                       PtxPredicate:$predicate,
                       UnitAttr:$tensormap,
                       UnitAttr:$uniform,
                       UnitAttr:$in_param_space);
  let assemblyFormat = "(`level` `=` $cacheLevel^ (`uniform` $uniform^)? `,`)? (`tensormap` $tensormap^ (`in_param_space` $in_param_space^)? `,`)? (`evict_priority` `=` $evictPriority^ `,`)? $addr (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static NVVM::IDArgPair
    getIntrinsicIDAndArgs(NVVM::PrefetchOp &op,LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
    bool hasIntrinsic() { return !getPredicate() || !getTensormap(); }
  }];
  let extraClassDefinition = [{
    std::string $cppClass::getPtx() {
      // Inline PTX is only supported for prefetch tensormap
      return std::string("prefetch.tensormap [%0];");
    }
  }];
  let llvmBuilder = [{
    auto [id, args] = NVVM::PrefetchOp::getIntrinsicIDAndArgs(op,
                                          moduleTranslation, builder);
    createIntrinsicCall(builder, id, builder.getVoidTy(), args);
  }];
}

def NVVM_CpAsyncBulkPrefetchOp : NVVM_VoidIntrinsicOp<"cp.async.bulk.prefetch"> {
  let summary = "Async bulk prefetch from global memory to L2 cache";
  let description = [{
    Initiates an asynchronous prefetch of data from the location
    specified by `srcMem` to the L2 cache.

    The `l2CacheHint` operand is optional, and it is used to specify cache
    eviction policy that may be used during the memory access.

    Example:
    ```mlir
      nvvm.cp.async.bulk.prefetch %src, %size : !llvm.ptr<1>

      // with l2_cache_hint
      nvvm.cp.async.bulk.prefetch %src, %size l2_cache_hint = %ch : !llvm.ptr<1>
    ```

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-prefetch)
  }];

  let arguments = (ins
    LLVM_PointerGlobal:$srcMem,
    I32:$size,
    Optional<I64>:$l2CacheHint);

  let assemblyFormat = [{
    $srcMem `,` $size (`l2_cache_hint` `=` $l2CacheHint^ )?
    attr-dict  `:` type($srcMem)
  }];
}

def NVVM_CpAsyncBulkTensorPrefetchOp :
  NVVM_VoidIntrinsicOp<"cp.async.bulk.tensor.prefetch", [AttrSizedOperandSegments]> {
  let arguments = (ins
    LLVM_PointerGeneric:$tmaDescriptor,
    Variadic<I32>:$coordinates,
    Variadic<I16>:$im2colOffsets,
    DefaultValuedAttr<TMALoadModeAttr, "TMALoadMode::TILE">:$mode,
    Optional<I64>:$l2CacheHint);

  let description = [{
    Initiates an asynchronous prefetch operation on the tensor data from global
    memory to L2 cache. This Op supports all the load modes specified in
    `TMALoadMode`.

    The `l2CacheHint` operand is optional, and it is used to specify cache
    eviction policy that may be used during the memory access.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-prefetch-tensor)
  }];

  let assemblyFormat = [{
    $tmaDescriptor `,`
    `box` `[`$coordinates `]`
    (`im2col` `[` $im2colOffsets^ `]` )?
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    attr-dict  `:` type($tmaDescriptor)
  }];

  let hasVerifier = 1;
}

// List of Reduction Ops supported with TMA Store
def TMAReduxKindAdd : I32EnumAttrCase<"ADD", 0, "add">;
def TMAReduxKindMin : I32EnumAttrCase<"MIN", 1, "min">;
def TMAReduxKindMax : I32EnumAttrCase<"MAX", 2, "max">;
def TMAReduxKindInc : I32EnumAttrCase<"INC", 3, "inc">;
def TMAReduxKindDec : I32EnumAttrCase<"DEC", 4, "dec">;
def TMAReduxKindAnd : I32EnumAttrCase<"AND", 5, "and">;
def TMAReduxKindOr  : I32EnumAttrCase<"OR",  6, "or">;
def TMAReduxKindXor : I32EnumAttrCase<"XOR", 7, "xor">;

def TMAReduxKind : I32EnumAttr<"TMAReduxKind", "NVVM TMA redux kind",
    [TMAReduxKindAdd, TMAReduxKindMax, TMAReduxKindMin,
     TMAReduxKindInc, TMAReduxKindDec, TMAReduxKindAnd,
     TMAReduxKindOr,  TMAReduxKindXor]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def TMAReduxKindAttr : EnumAttr<NVVM_Dialect, TMAReduxKind, "tma_redux_kind"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_CpAsyncBulkTensorReduceOp :
  NVVM_Op<"cp.async.bulk.tensor.reduce", [AttrSizedOperandSegments]> {
  let arguments = (ins
    LLVM_AnyPointer:$tmaDescriptor,
    LLVM_PointerShared:$srcMem,
    TMAReduxKindAttr:$redKind,
    DefaultValuedAttr<TMAStoreModeAttr, "TMAStoreMode::TILE">:$mode,
    Variadic<I32>:$coordinates,
    Optional<I64>:$l2CacheHint);

  let description = [{
    Initiates an asynchronous reduction operation of tensor data in
    global memory with tensor data in shared memory.

    The `mode` attribute indicates whether the copy mode is tile or im2col.
    The `redOp` attribute specifies the reduction operations applied.
    The supported reduction operations are:
    {add, min, max, inc, dec, and, or, xor}

    The `l2CacheHint` operand is optional, and it is used to specify cache
    eviction policy that may be used during the memory access.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor)
  }];

  let assemblyFormat = [{
    $tmaDescriptor `,`
    $srcMem `,`
    `box` `[`$coordinates `]`
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    attr-dict  `:` type($tmaDescriptor) `,` type($srcMem)
  }];

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op,
      LLVM::ModuleTranslation &mt, llvm::IRBuilderBase& builder);
  }];

  let hasVerifier = 1;

  string llvmBuilder = [{
    auto [id, args] = NVVM::CpAsyncBulkTensorReduceOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_CpAsyncBulkGlobalToSharedClusterOp :
  NVVM_Op<"cp.async.bulk.shared.cluster.global", [AttrSizedOperandSegments]> {
  let summary = "Async bulk copy from global to Shared {cta or cluster} memory";
  let description = [{
    Initiates an asynchronous copy operation from global memory to shared
    memory or shared_cluster memory.

    The `multicastMask` operand is optional and can be used only when the
    destination is shared::cluster memory. When it is present, this Op copies
    data from global memory to shared memory of multiple CTAs in the cluster.
    Operand `multicastMask` specifies the destination CTAs in the cluster such
    that each bit position in the 16-bit `multicastMask` operand corresponds to
    the `nvvm.read.ptx.sreg.ctaid` of the destination CTA. 

    The `l2CacheHint` operand is optional, and it is used to specify cache
    eviction policy that may be used during the memory access.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_PointerShared, LLVM_PointerSharedCluster]>:$dstMem,
    LLVM_PointerGlobal:$srcMem,
    LLVM_PointerShared:$mbar,
    I32:$size,
    Optional<I16>:$multicastMask,
    Optional<I64>:$l2CacheHint);

  let assemblyFormat = [{
    $dstMem `,` $srcMem `,` $mbar `,` $size
    (`multicast_mask` `=` $multicastMask^ )?
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    attr-dict  `:` type($dstMem) `,` type($srcMem)
  }];

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];
  string llvmBuilder = [{
    auto [id, args] = NVVM::CpAsyncBulkGlobalToSharedClusterOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_CpAsyncBulkSharedCTAToSharedClusterOp :
  NVVM_Op<"cp.async.bulk.shared.cluster.shared.cta"> {
  let summary = "Async bulk copy from Shared CTA memory to Shared cluster memory";
  let description = [{
    Initiates an asynchronous copy operation from Shared CTA memory to Shared
    cluster memory.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk)
  }];

  let arguments = (ins
    LLVM_PointerSharedCluster:$dstMem,
    LLVM_PointerShared:$srcMem,
    LLVM_PointerShared:$mbar,
    I32:$size);

  let assemblyFormat = [{
    $dstMem `,` $srcMem `,` $mbar `,` $size
    attr-dict  `:` type($dstMem) `,` type($srcMem)
  }];

  string llvmBuilder = [{
    createIntrinsicCall(builder,
      llvm::Intrinsic::nvvm_cp_async_bulk_shared_cta_to_cluster,
      {$dstMem, $mbar, $srcMem, $size});
  }];
}

def NVVM_CpAsyncBulkSharedCTAToGlobalOp :
  NVVM_Op<"cp.async.bulk.global.shared.cta", [AttrSizedOperandSegments]> {
  let summary = "Async bulk copy from Shared CTA memory to Global memory";
  let description = [{
    Initiates an asynchronous copy operation from Shared CTA memory to
    global memory. The 32-bit operand `size` specifies the amount of
    memory to be copied, in terms of number of bytes. `size` must be a
    multiple of 16. The `l2CacheHint` operand is optional, and it is used
    to specify cache eviction policy that may be used during the memory
    access. The `byteMask` operand is optional. The i-th bit in the 16-bit
    wide `byteMask` specifies whether the i-th byte of each 16-byte wide
    chunk of source data is copied to the destination. If the bit is set,
    the byte is copied.

    Example:
    ```mlir
      nvvm.cp.async.bulk.global.shared.cta %dst, %src, %size
          : !llvm.ptr<1>, !llvm.ptr<3>

      // with l2_cache_hint
      nvvm.cp.async.bulk.global.shared.cta %dst, %src, %size l2_cache_hint = %ch
          : !llvm.ptr<1>, !llvm.ptr<3>

      // with byte_mask
      nvvm.cp.async.bulk.global.shared.cta %dst, %src, %size byte_mask = %mask
          : !llvm.ptr<1>, !llvm.ptr<3>

      // with both l2_cache_hint and byte_mask
      nvvm.cp.async.bulk.global.shared.cta %dst, %src, %size l2_cache_hint = %ch byte_mask = %mask
          : !llvm.ptr<1>, !llvm.ptr<3>
    ```

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk)
  }];

  let arguments = (ins
    LLVM_PointerGlobal:$dstMem,
    LLVM_PointerShared:$srcMem,
    I32:$size,
    Optional<I64>:$l2CacheHint,
    Optional<I16>:$byteMask);

  let assemblyFormat = [{
    $dstMem `,` $srcMem `,` $size
    (`l2_cache_hint` `=` $l2CacheHint^ )?
    (`byte_mask` `=` $byteMask^ )?
    attr-dict `:` type($dstMem) `,` type($srcMem)
  }];

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase& builder);
  }];
  string llvmBuilder = [{
    auto [id, args] = NVVM::CpAsyncBulkSharedCTAToGlobalOp::getIntrinsicIDAndArgs(
                      *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM Wgmma Ops
//===----------------------------------------------------------------------===//

def NVVM_WgmmaFenceAlignedOp : NVVM_Op<"wgmma.fence.aligned", [NVVMRequiresSMa<[90]>]> {
  let arguments = (ins);
  let description = [{
    Enforce an ordering of register accesses between warpgroup level matrix 
    multiplication and other operations. 
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions-wgmma-fence)
  }];
  let assemblyFormat = "attr-dict";
  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_wgmma_fence_sync_aligned);
  }];
}

def NVVM_WgmmaGroupSyncAlignedOp : NVVM_Op<"wgmma.commit.group.sync.aligned", [NVVMRequiresSMa<[90]>]> {
  let assemblyFormat = "attr-dict";
  let description = [{
    Commits all prior uncommitted warpgroup level matrix multiplication operations.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions-wgmma-commit-group)
  }];
  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_wgmma_commit_group_sync_aligned);
  }];
}

def NVVM_WgmmaWaitGroupSyncOp : NVVM_Op<"wgmma.wait.group.sync.aligned", [NVVMRequiresSMa<[90]>]> {
  let arguments = (ins I64Attr:$group);
  let assemblyFormat = "attr-dict $group";
  let description = [{
    Signal the completion of a preceding warpgroup operation.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions-wgmma-wait-group)
  }];
  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_wgmma_wait_group_sync_aligned, builder.getInt64($group));
  }];
}

/// Enum attribute type for the negating of input operands
def WGMMAScaleInNeg : I32EnumAttrCase<"neg", -1>;
def WGMMAScaleInOne : I32EnumAttrCase<"one", 1>;
def WGMMAScaleIn : I32EnumAttr<"WGMMAScaleIn", "WGMMA overflow options",
  [WGMMAScaleInOne, WGMMAScaleInNeg]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def WGMMAScaleInAttr : EnumAttr<NVVM_Dialect, WGMMAScaleIn, "wgmma_scale_in"> {
  let assemblyFormat = "`<` $value `>`";
}

/// Enum attribute type for the output operand
def WGMMAScaleOutZero : I32EnumAttrCase<"zero", 0>;
def WGMMAScaleOutOne : I32EnumAttrCase<"one", 1>;
def WGMMAScaleOut : I32EnumAttr<"WGMMAScaleOut", "WGMMA input predicate",
  [WGMMAScaleOutZero, WGMMAScaleOutOne]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def WGMMAScaleOutAttr : EnumAttr<NVVM_Dialect, WGMMAScaleOut, "wgmma_scale_out"> {
  let assemblyFormat = "`<` $value `>`";
}

/// Enum attribute of the different PTX element types used for WGMMA operands.
def WGMMATypeF16  : I32EnumAttrCase<"f16", 0>;
def WGMMATypeTF32 : I32EnumAttrCase<"tf32", 1>;
def WGMMATypeU8 : I32EnumAttrCase<"u8", 2>;
def WGMMATypeS8 : I32EnumAttrCase<"s8", 3>;
def WGMMATypeB1 : I32EnumAttrCase<"b1", 4>;
def WGMMATypeBF16 : I32EnumAttrCase<"bf16", 5>;
def WGMMATypeF8E4M3 : I32EnumAttrCase<"e4m3", 6>;
def WGMMATypeF8E5M2 : I32EnumAttrCase<"e5m2", 7>;
def WGMMATypeF32 : I32EnumAttrCase<"f32", 8>;
def WGMMATypeS32 : I32EnumAttrCase<"s32", 9>;

def WGMMATypes : I32EnumAttr<"WGMMATypes", "NVVM WGMMA types",
  [WGMMATypeF16, WGMMATypeTF32,
    WGMMATypeU8, WGMMATypeS8,
    WGMMATypeB1, WGMMATypeBF16, WGMMATypeF8E4M3, 
    WGMMATypeF8E5M2, WGMMATypeF32, WGMMATypeS32]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def WGMMATypesAttr : EnumAttr<NVVM_Dialect, WGMMATypes, "wgmma_type"> {
  let assemblyFormat = "`<` $value `>`";
}


def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async", 
              [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>,
                PredOpTrait<"input struct and result struct must be the same type",
                  TCresIsSameAsOpBase<0, 0>>,]> 
{
  let results = (outs LLVM_AnyStruct:$results);
  let arguments = (ins 
    LLVM_AnyStruct:$inouts,
    I64:$descriptorA, 
    I64:$descriptorB, 
    NVVM_MMAShapeAttr:$shape,
    WGMMATypesAttr:$typeA,
    WGMMATypesAttr:$typeB,
    WGMMATypesAttr:$typeD,
    WGMMAScaleOutAttr:$scaleD,
    WGMMAScaleInAttr:$scaleA,
    WGMMAScaleInAttr:$scaleB, 
    MMALayoutAttr:$layoutA,
    MMALayoutAttr:$layoutB,
    OptionalAttr<MMAIntOverflowAttr>:$satfinite
  );  
  
   let assemblyFormat = [{ 
      $descriptorA `,` $descriptorB `,` $inouts `,` $shape `,`
      `D` `[` $typeD `,` $scaleD (`,` $satfinite^)? `]` `,`
      `A` `[` $typeA `,` $scaleA `,` $layoutA `]` `,` 
      `B` `[` $typeB `,` $scaleB `,` $layoutB `]`
      attr-dict `:` 
      type($inouts) `->` type($results)
    }];
  
  let description = [{
    The warpgroup (128 threads) level matrix multiply and accumulate operation 
    has either of the following forms, where matrix D is called accumulator:
      D = A * B + D
      D = A * B, where the input from accumulator D is disabled.

    Supported shapes:  
    ```
    |--------------|--------------|------------|--------------|---------------|
    |              |              |            |              |f16+=e4m3*e4m3 |
    |              |              |            |              |f16+=e5m2*e5m2 |
    |f32+=tf32*tf32|f16+=f16 *f16 | s32+=s8*s8 |s32 += b1 * b1|f16+=e5m2*e4m3 |
    |              |f32+=f16 *f16 | s32+=u8*u8 |              |f16+=e4m3*e5m2 |
    |              |f32+=bf16*bf16| s32+=u8*u8 |              |f16+=e4m3*e5m2 |
    |              |f32+=bf16*bf16| s32+=s8*u8 |              |f32+=e4m3*e4m3 |
    |              |              | s32+=u8*s8 |              |f32+=e5m2*e5m2 |
    |              |              |            |              |f32+=e4m3*e5m2 |
    |              |              |            |              |f32+=e4m3*e5m2 |
    |--------------|--------------|------------|--------------|---------------|
    |   .m64n8k8   |  .m64n8k16   | .m64n8k32  | .m64n8k256   | .m64n8k32     |
    |   .m64n16k8  |  .m64n16k16  | .m64n16k32 | .m64n16k256  | .m64n16k32    |
    |   .m64n24k8  |  .m64n24k16  | .m64n24k32 | .m64n24k256  | .m64n24k32    |
    |   .m64n32k8  |  .m64n32k16  | .m64n32k32 | .m64n32k256  | .m64n32k32    |
    |   .m64n40k8  |  .m64n40k16  | .m64n48k32 | .m64n48k256  | .m64n40k32    |
    |   .m64n48k8  |  .m64n48k16  | .m64n64k32 | .m64n64k256  | .m64n48k32    |
    |   .m64n56k8  |  .m64n56k16  | .m64n80k32 | .m64n80k256  | .m64n56k32    |
    |   .m64n64k8  |  .m64n64k16  | .m64n96k32 | .m64n96k256  | .m64n64k32    |
    |   .m64n72k8  |  .m64n72k16  | .m64n112k32| .m64n112k256 | .m64n72k32    |
    |   .m64n80k8  |  .m64n80k16  | .m64n128k32| .m64n128k256 | .m64n80k32    |
    |   .m64n88k8  |  .m64n88k16  | .m64n144k32| .m64n144k256 | .m64n88k32    |
    |   .m64n96k8  |  .m64n96k16  | .m64n160k32| .m64n160k256 | .m64n96k32    |
    |   .m64n104k8 |  .m64n104k16 | .m64n176k32| .m64n176k256 | .m64n104k32   |
    |   .m64n112k8 |  .m64n112k16 | .m64n192k32| .m64n192k256 | .m64n112k32   |
    |   .m64n120k8 |  .m64n120k16 | .m64n208k32| .m64n208k256 | .m64n120k32   |
    |   .m64n128k8 |  .m64n128k16 | .m64n224k32| .m64n224k256 | .m64n128k32   |
    |   .m64n136k8 |  .m64n136k16 | .m64n240k32| .m64n240k256 | .m64n136k32   |
    |   .m64n144k8 |  .m64n144k16 | .m64n256k32| .m64n256k256 | .m64n144k32   |
    |   .m64n152k8 |  .m64n152k16 |            |              | .m64n152k32   |
    |   .m64n160k8 |  .m64n160k16 |            |              | .m64n160k32   |
    |   .m64n168k8 |  .m64n168k16 |            |              | .m64n168k32   |
    |   .m64n176k8 |  .m64n176k16 |            |              | .m64n176k32   |
    |   .m64n184k8 |  .m64n184k16 |            |              | .m64n184k32   |
    |   .m64n192k8 |  .m64n192k16 |            |              | .m64n192k32   |
    |   .m64n200k8 |  .m64n200k16 |            |              | .m64n200k32   |
    |   .m64n208k8 |  .m64n208k16 |            |              | .m64n208k32   |
    |   .m64n216k8 |  .m64n216k16 |            |              | .m64n216k32   |
    |   .m64n224k8 |  .m64n224k16 |            |              | .m64n224k32   |
    |   .m64n232k8 |  .m64n232k16 |            |              | .m64n232k32   |
    |   .m64n240k8 |  .m64n240k16 |            |              | .m64n240k32   |
    |   .m64n248k8 |  .m64n248k16 |            |              | .m64n248k32   |
    |   .m64n256k8 |  .m64n256k16 |            |              | .m64n256k32   |
    |--------------|--------------|------------|--------------|---------------|
    ```

    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions)
  }];
  
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    bool getAsmValues(RewriterBase &, llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM Griddepcontrol Ops
//===----------------------------------------------------------------------===//

def GridDepActionWait : I32EnumCase<"wait", 0>;
def GridDepActionLaunchDependent : I32EnumCase<"launch_dependents", 1>;

def GridDepActionKind : I32Enum<"GridDepActionKind", "Action kind for grid dependency control",
  [GridDepActionWait, GridDepActionLaunchDependent]> {
  let cppNamespace = "::mlir::NVVM";
}

def GridDepActionAttr : EnumAttr<NVVM_Dialect, GridDepActionKind, "grid_dep_action">;

def NVVM_GriddepcontrolOp : NVVM_Op<"griddepcontrol", []> {
  let description = [{
    If the $kind attribute is set to `wait`, it causes the 
    executing thread to wait until all prerequisite grids in flight 
    have completed and all the memory operations from the prerequisite grids 
    are performed and made visible to the current grid.

    When the $kind is launch_dependents, it signals that specific dependents 
    the runtime system designated to react to this instruction can be scheduled 
    as soon as all other CTAs in the grid issue the same instruction or have 
    completed.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-griddepcontrol)
  }];

  let arguments = (ins GridDepActionAttr:$kind);

  let assemblyFormat = "$kind attr-dict";

  string llvmBuilder = [{
    llvm::Intrinsic::ID id;
      switch ($kind) {
        case NVVM::GridDepActionKind::wait:
          id = llvm::Intrinsic::nvvm_griddepcontrol_wait;
          break;
        case NVVM::GridDepActionKind::launch_dependents:
          id = llvm::Intrinsic::nvvm_griddepcontrol_launch_dependents;
          break;
      }
      createIntrinsicCall(builder, id);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM Mapa Op
//===----------------------------------------------------------------------===//

def NVVM_MapaOp: NVVM_Op<"mapa",
  [InputAddressIsCombinationOf<["a", "res"],
    [[LLVM_PointerShared, LLVM_PointerSharedCluster], [LLVM_PointerGeneric, LLVM_PointerGeneric]],
    "Valid address-space check(or mapping) for mapa Op">, NVVMRequiresSM<90>]> {
  let results = (outs AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerSharedCluster]>:$res);
  let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$a, I32:$b);

  string llvmBuilder = [{
    int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
    
    bool isSharedMemory = addrSpace == static_cast<int> (NVVM::NVVMMemorySpace::Shared);

    auto intId = isSharedMemory? llvm::Intrinsic::nvvm_mapa_shared_cluster : llvm::Intrinsic::nvvm_mapa;
    $res = createIntrinsicCall(builder, intId, {$a, $b});
  }];
  
  let assemblyFormat = "$a`,` $b attr-dict `:` type($a) `->` type($res)";
}

//===----------------------------------------------------------------------===//
// NVVM match.sync Op
//===----------------------------------------------------------------------===//

def MatchSyncKindAny : I32EnumAttrCase<"any", 0>;
def MatchSyncKindAll : I32EnumAttrCase<"all", 1>;

def MatchSyncKind : I32EnumAttr<"MatchSyncKind", "NVVM match sync kind",
  [MatchSyncKindAny, MatchSyncKindAll]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}

def MatchSyncKindAttr : EnumAttr<NVVM_Dialect, MatchSyncKind, "match_sync_kind">;

def NVVM_MatchSyncOp : NVVM_Op<"match.sync", [InferTypeOpAdaptor]>,
  Results<(outs AnyTypeOf<[I32, LLVMStructType]>:$res)>,
  Arguments<(ins I32:$thread_mask,
                 AnyTypeOf<[I32, I64]>:$val,
                 MatchSyncKindAttr:$kind)> {
  let summary = "Broadcast and compare a value across threads in warp";
  let description = [{
    The `match.sync` op performs broadcast and compare of operand `val` across 
    all non-exited threads in `thread_mask` and returns a mask depending on the 
    kind and an optional predicate.

    The matching operation kinds are:
    - `any`: Returns a mask corresponding to the non-exited threads in the 
    `thread_mask` that have the same value of operand `val`.
    - `all`: Returns a mask and a predicate. If all non-exited threads in the 
    `thread_mask` have the same value of operand `val`, the predicate is set to 
    true and the mask corresponds to the non-exited threads in the 
    `thread_mask`. Otherwise, the predicate is set to false and the mask is 0.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-match-sync)
  }];
  string llvmBuilder = [{
    auto intId = getMatchSyncIntrinsicId(
        op.getVal().getType(), $kind);
    $res = createIntrinsicCall(builder,
        intId, {$thread_mask, $val});
  }];
  let assemblyFormat = "$kind $thread_mask `,` $val attr-dict `:` type($val) `->` type($res)";
  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// NVVM Bulk Store Op
//===----------------------------------------------------------------------===//

def NVVM_BulkStoreOp: NVVM_Op<"st.bulk"> {
  let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$addr, I64:$size, DefaultValuedAttr<I64Attr, "0">:$initVal);

  let summary = "Bulk Store Op";
  let description = [{
    Initializes a region of shared memory at the address given by `addr`.
    The `size` operand specifies the number of bytes to initialize and must be 
    a multiple of 8.
    The `initVal` operand specifies the value to initialize the memory to. The 
    only supported value is 0.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-st-bulk)
  }];

  string llvmBuilder = [{
    auto intId = getStBulkIntrinsicId(
          llvm::cast<LLVM::LLVMPointerType>(op.getAddr().getType()));
    createIntrinsicCall(builder, intId,
                      {$addr, $size, builder.getInt64($initVal)});
  }];
  
  let assemblyFormat = "$addr `,` `size` `=` $size (`,` `init` `=` $initVal^)? attr-dict `:` type($addr)";
  
  let hasVerifier = 1;
}

def NVVM_Exit : NVVM_Op<"exit"> {
  let summary = "Exit Op";
  let description = [{
    Ends execution of a thread.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#control-flow-instructions-exit)
  }];
  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_exit);
  }];

  let assemblyFormat = "attr-dict";
}


//===----------------------------------------------------------------------===//
// NVVM breakpoint Op
//===----------------------------------------------------------------------===//

def NVVM_Breakpoint : NVVM_Op<"breakpoint"> {
  let summary = "Breakpoint Op";
  let description = [{
    Breakpoint suspends execution of the program for debugging.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-brkpt)
  }];
  string llvmBuilder = [{
    createIntrinsicCall(builder, llvm::Intrinsic::debugtrap);
  }];

  let assemblyFormat = "attr-dict";
}

//===----------------------------------------------------------------------===//
// NVVM TCGEN05 Ops
//===----------------------------------------------------------------------===//
def Tcgen05FenceBefore : I32EnumAttrCase<"BEFORE_THREAD_SYNC", 0, "before">;
def Tcgen05FenceAfter  : I32EnumAttrCase<"AFTER_THREAD_SYNC",  1, "after">;
def Tcgen05FenceKind : I32EnumAttr<"Tcgen05FenceKind", "NVVM Tcgen05 fence kind",
  [Tcgen05FenceBefore, Tcgen05FenceAfter]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def Tcgen05FenceKindAttr :
  EnumAttr<NVVM_Dialect, Tcgen05FenceKind, "tcgen05_fence"> {
  let assemblyFormat = "`<` $value `>`";
}

def Tcgen05WaitLoad  : I32EnumAttrCase<"LOAD",  0, "load">;
def Tcgen05WaitStore : I32EnumAttrCase<"STORE", 1, "store">;
def Tcgen05WaitKind : I32EnumAttr<"Tcgen05WaitKind", "NVVM Tcgen05 wait kind",
  [Tcgen05WaitLoad, Tcgen05WaitStore]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::mlir::NVVM";
}
def Tcgen05WaitKindAttr :
  EnumAttr<NVVM_Dialect, Tcgen05WaitKind, "tcgen05_wait"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_Tcgen05AllocOp : NVVM_Op<"tcgen05.alloc", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 alloc operation";
  let description = [{
    The `tcgen05.alloc` Op allocates tensor core memory for
    the amount specified by `nCols` and writes the destination
    address to the `addr` argument. The `nCols` operand specifies the
    number of columns to be allocated and it must be a power-of-two.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-alloc-manage-instructions)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_AnyPointer, LLVM_PointerShared]>:$addr,
    I32:$nCols,
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);

  let assemblyFormat = "$addr `,` $nCols attr-dict `:` type(operands)";

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::SmallVector<llvm::Value *> &args);
  }];
  string llvmBuilder = [{
    llvm::SmallVector<llvm::Value *> args;
    auto id = NVVM::Tcgen05AllocOp::getIntrinsicIDAndArgs(
      *op, moduleTranslation, args);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_Tcgen05DeallocOp : NVVM_Op<"tcgen05.dealloc", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 dealloc operation";
  let description = [{
    The `tcgen05.dealloc` Op de-allocates the tensor core memory
    specified by `tmemAddr`, which must be from a previous tensor
    memory allocation. The `nCols` operand specifies the number
    of columns to be de-allocated, and it must be a power-of-two.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-alloc-manage-instructions)
  }];

  let arguments = (ins LLVM_PointerTensor:$taddr, I32:$nCols,
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);

  let assemblyFormat = "$taddr `,` $nCols attr-dict `:` type(operands)";

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::SmallVector<llvm::Value *> &args);
  }];
  string llvmBuilder = [{
    llvm::SmallVector<llvm::Value *> args;
    auto id = NVVM::Tcgen05DeallocOp::getIntrinsicIDAndArgs(
      *op, moduleTranslation, args);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_Tcgen05RelinquishAllocPermitOp : NVVM_Op<"tcgen05.relinquish_alloc_permit",
    [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 Op to relinquish the right to allocate";
  let description = [{
    The `tcgen05.relinquish_alloc_permit` Op specifies that the CTA
    of the executing thread is relinquishing the right to allocate
    Tensor Memory. So, it is illegal for a CTA to perform `tcgen05.alloc`
    after any of its constituent threads execute `tcgen05.relinquish_alloc_permit`.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-alloc-manage-instructions)
  }];

  let arguments = (ins
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);

  let assemblyFormat = "attr-dict";

  string llvmBuilder = [{
    auto id = ($group == NVVM::CTAGroupKind::CTA_1) ?
      llvm::Intrinsic::nvvm_tcgen05_relinq_alloc_permit_cg1 :
      llvm::Intrinsic::nvvm_tcgen05_relinq_alloc_permit_cg2;
    createIntrinsicCall(builder, id);
  }];
}

def NVVM_Tcgen05FenceOp : NVVM_Op<"tcgen05.fence", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 fence operations";
  let description = [{
    The `tcgen05.fence<before>` orders all prior async tcgen05 operations
    with respect to the subsequent tcgen05 and execution ordering operations.
    The `tcgen05.fence<after>` orders all subsequent async tcgen05 operations
    with respect to the prior tcgen05 and execution ordering operations.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensorcore-5th-generation-instructions-tcgen05-fence)
  }];

  let arguments = (ins Tcgen05FenceKindAttr:$kind);
  let assemblyFormat = "$kind attr-dict";

  string llvmBuilder = [{
    auto id = ($kind == NVVM::Tcgen05FenceKind::BEFORE_THREAD_SYNC)
      ? llvm::Intrinsic::nvvm_tcgen05_fence_before_thread_sync
      : llvm::Intrinsic::nvvm_tcgen05_fence_after_thread_sync;
    createIntrinsicCall(builder, id);
  }];
}

def NVVM_Tcgen05WaitOp : NVVM_Op<"tcgen05.wait", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 wait operations";
  let description = [{
    The `tcgen05.wait<load>` causes the executing thread to block until
    all prior `tcgen05.ld` operations issued by the executing thread
    have completed. Similarly, the `tcgen05.wait<store>` causes the executing
    thread to block until all prior `tcgen05.st` operations issued by the
    executing thread have completed.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-wait)
  }];

  let arguments = (ins Tcgen05WaitKindAttr:$kind);
  let assemblyFormat = "$kind attr-dict";

  string llvmBuilder = [{
    auto id = ($kind == NVVM::Tcgen05WaitKind::LOAD)
      ? llvm::Intrinsic::nvvm_tcgen05_wait_ld
      : llvm::Intrinsic::nvvm_tcgen05_wait_st;
    createIntrinsicCall(builder, id);
  }];
}

def NVVM_Tcgen05CommitOp : NVVM_Op<"tcgen05.commit", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 commit operations";
  let description = [{
    The `tcgen05.commit` makes the *mbarrier object*, specified by
    the operand `addr`, track the completion of all the prior
    async-tcgen05 operations initiated by the executing thread.
    The multicast variants allow signaling on the *mbarrier objects*
    of multiple CTAs within the cluster. Operand `multicastMask`,
    when present, specifies the destination CTAs in the cluster such
    that each bit position in the 16-bit `multicastMask` operand
    corresponds to the `nvvm.read.ptx.sreg.ctaid` of the destination CTA.
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen-async-sync-operations-commit)
  }];

  let arguments = (ins
    AnyTypeOf<[LLVM_AnyPointer, LLVM_PointerShared]>:$addr,
    Optional<I16>:$multicastMask,
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);

  let assemblyFormat = [{
    $addr (`,` `multicast_mask` `=` $multicastMask^)?
    attr-dict `:` type(operands)
  }];

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID
      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                            llvm::SmallVector<llvm::Value *> &args);
  }];

  string llvmBuilder = [{
    llvm::SmallVector<llvm::Value *> args;
    auto id = NVVM::Tcgen05CommitOp::getIntrinsicIDAndArgs(
      *op, moduleTranslation, args);
    createIntrinsicCall(builder, id, args);
  }];
}

def NVVM_Tcgen05ShiftOp : NVVM_Op<"tcgen05.shift",
    [NVVMRequiresSMa<[100, 101, 103, 110]>]> {
  let summary = "Tcgen05 shift operation";
  let description = [{
    The `tcgen05.shift` is an asynchronous instruction which initiates
    the shifting of 32-byte elements downwards across all the rows,
    except the last, by one row. The operand `taddr` specifies the base
    address of the matrix in Tensor Memory whose rows must be down shifted.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-shift)
  }];

  let arguments = (ins LLVM_PointerTensor:$taddr,
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);

  let assemblyFormat = "$taddr attr-dict `:` type(operands)";

  string llvmBuilder = [{
    auto id = ($group == NVVM::CTAGroupKind::CTA_1) ?
      llvm::Intrinsic::nvvm_tcgen05_shift_down_cg1 :
      llvm::Intrinsic::nvvm_tcgen05_shift_down_cg2;
    createIntrinsicCall(builder, id, {$taddr});
  }];
}

def Shape128x256b : I32EnumAttrCase<"SHAPE_128x256b", 0, "shape_128x256b">;
def Shape4x256b   : I32EnumAttrCase<"SHAPE_4x256b",   1, "shape_4x256b">;
def Shape128x128b : I32EnumAttrCase<"SHAPE_128x128b", 2, "shape_128x128b">;
def Shape64x128b  : I32EnumAttrCase<"SHAPE_64x128b",  3, "shape_64x128b">;
def Shape32x128b  : I32EnumAttrCase<"SHAPE_32x128b",  4, "shape_32x128b">;

def Tcgen05CpShape : I32EnumAttr<"Tcgen05CpShape", "tcgen05 cp shapes",
  [Shape128x256b, Shape4x256b, Shape128x128b, Shape64x128b, Shape32x128b]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}
def Tcgen05CpShapeAttr : EnumAttr<NVVM_Dialect, Tcgen05CpShape, "tcgen05_cp_shape"> {
  let assemblyFormat = "`<` $value `>`";
}

def Tcgen05CpMulticastNone: I32EnumAttrCase<"NONE", 0, "none">;
def Tcgen05CpMulticastWarpx2_02_13: I32EnumAttrCase<"WARPX2_02_13", 1, "warpx2_02_13">;
def Tcgen05CpMulticastWarpx2_01_23: I32EnumAttrCase<"WARPX2_01_23", 2, "warpx2_01_23">;
def Tcgen05CpMulticastWarpx4: I32EnumAttrCase<"WARPX4", 3, "warpx4">;

def Tcgen05CpMulticast : I32EnumAttr<"Tcgen05CpMulticast", "tcgen05 cp multicast",
  [Tcgen05CpMulticastNone, Tcgen05CpMulticastWarpx2_02_13,
   Tcgen05CpMulticastWarpx2_01_23, Tcgen05CpMulticastWarpx4]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}
def Tcgen05CpMulticastAttr : EnumAttr<NVVM_Dialect, Tcgen05CpMulticast, "tcgen05_cp_multicast"> {
  let assemblyFormat = "`<` $value `>`";
}

def FormatB6x16_P32: I32EnumAttrCase<"B6x16_P32", 0, "b6x16_p32">;
def FormatB4x16_P64: I32EnumAttrCase<"B4x16_P64", 1, "b4x16_p64">;

def Tcgen05CpSrcFormat : I32EnumAttr<"Tcgen05CpSrcFormat", "tcgen05 cp source format",
  [FormatB6x16_P32, FormatB4x16_P64]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}
def Tcgen05CpSrcFormatAttr : EnumAttr<NVVM_Dialect, Tcgen05CpSrcFormat, "tcgen05_cp_src_fmt"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_Tcgen05CpOp : NVVM_Op<"tcgen05.cp", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Tcgen05 copy operation";
  let description = [{
    Instruction tcgen05.cp initiates an asynchronous copy operation from
    shared memory to the location specified by the address operand `taddr`
    in the Tensor Memory. The 64-bit register operand `smem_desc` specifies
    the matrix descriptor representing the source matrix in the shared memory
    that needs to be copied.

    Example:
    ```mlir
      nvvm.tcgen05.cp %taddr, %smem_desc {
        group = #nvvm.tcgen05_group<cta_2>,
        shape = #nvvm.tcgen05_cp_shape<shape_64x128b>,
        multicast = #nvvm.tcgen05_cp_multicast<warpx2_01_23>,
        srcFormat = #nvvm.tcgen05_cp_src_fmt<b6x16_p32>
      }
    ```
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensorcore-5th-generation-instructions-tcgen05-cp)
  }];

  let arguments = (ins
    Tcgen05CpShapeAttr:$shape,
    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group,
    DefaultValuedAttr<Tcgen05CpMulticastAttr, "Tcgen05CpMulticast::NONE">:$multicast,
    OptionalAttr<Tcgen05CpSrcFormatAttr>:$srcFormat,
    LLVM_PointerTensor:$taddr,
    I64:$smem_desc);

  let assemblyFormat = "$taddr`,` $smem_desc attr-dict";
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static llvm::Intrinsic::ID getIntrinsicID(Operation &op);
  }];

  string llvmBuilder = [{
    auto id = NVVM::Tcgen05CpOp::getIntrinsicID(*op);
    createIntrinsicCall(builder, id, {$taddr, $smem_desc});
  }];
}

def NVVM_Tcgen05MmaSmemDescOp : NVVM_Op<"tcgen05.mma_smem_desc", []> {
  let summary = "Constructs a Shared Memory descriptor for MMA Operands A or B";
  let description = [{
    The `nvvm.tcgen05_mma_smem_desc` constructs a Shared Memory descriptor
    for tcgen05.mma. This descriptor is a 64-bit value which describes the
    properties of multiplicand matrix in shared memory including its location
    in the shared memory of the current CTA.

    ```
    +-----------+------+------------------------------------------------------+
    | Bit-field | Size | Description                                          |
    +-----------+------+------------------------------------------------------+
    | 0-13      | 14   | Matrix start address                                 |
    | 14-15     | 2    | Reserved                                             |
    | 16-29     | 14   | Leading dim relative-offset (or) absolute-address    |
    | 30-31     | 2    | Reserved                                             |
    | 32-45     | 14   | Stride dimension byte offset                         |
    | 46-48     | 3    | Fixed constant value of 0b001                        |
    | 49-51     | 3    | Matrix base offset                                   |
    | 52        | 1    | Leading dimension stride mode:                       |
    |           |      |   0: byte offset relative                            |
    |           |      |   1: byte address absolute                           |
    | 53-60     | 8    | Fixed constant value of 0xb00000000                  |
    | 61-63     | 3    | Swizzling mode:                                      |
    |           |      |   0: No swizzling                                    |
    |           |      |   1: 128-Byte with 32B atomic swizzling              |
    |           |      |   2: 128-Byte swizzling                              |
    |           |      |   4: 64-Byte swizzling                               |
    |           |      |   6: 32-Byte swizzling                               |
    |           |      |   (Values 3, 5 and 7 are invalid)                    |
    +-----------+------+------------------------------------------------------+    
    ```

    Example:
    ```mlir
      %desc = nvvm.tcgen05.mma_smem_desc (%startAddr, %leadingDimOffset, %strideDimOffset,
                                          %baseOffset, %leadingDimMode, %swizzleMode) : (i32, i32, i32, i8, i1, i8) -> i64
    ```
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-shared-memory-descriptor)
  }];

  let arguments = (ins
      I32:$startAddr,        // Matrix A or B start address (bits 13-0)
      I32:$leadingDimOffset, // Matrix A or B leading dim byte offset (bits 29-16)
      I32:$strideDimOffset,  // Matrix A or B stride  dim byte offset (bits 45-32)
      I8:$baseOffset,        // Matrix A or B base offset (bits 51-49)
      I1:$leadingDimMode,    // Matrix A or B leading dim mode (bit 52)
      I8:$swizzleMode        // Swizzle mode (bits 63-61)
  );

  let results = (outs I64:$res);

  let assemblyFormat = [{
    `(` operands `)` attr-dict `:` `(` type(operands) `)` `->` type($res)
  }];

  let extraClassDeclaration = [{
    static void createSmemDescriptor(Operation &op, LLVM::ModuleTranslation &mt,
                                     llvm::IRBuilderBase& builder);
  }];

  string llvmBuilder = [{
    NVVM::Tcgen05MmaSmemDescOp::createSmemDescriptor(*op, moduleTranslation, builder);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM tcgen05 LdSt Shape Attr
//===----------------------------------------------------------------------===//

def Tcgen05LdStShape16x64b: I32EnumAttrCase<"SHAPE_16X64B", 0, "shape_16x64b">;
def Tcgen05LdStShape16x128b: I32EnumAttrCase<"SHAPE_16X128B", 1, "shape_16x128b">;
def Tcgen05LdStShape16x256b: I32EnumAttrCase<"SHAPE_16X256B", 2, "shape_16x256b">;
def Tcgen05LdStShape32x32b: I32EnumAttrCase<"SHAPE_32X32B", 3, "shape_32x32b">;
def Tcgen05LdStShape16x32bx2: I32EnumAttrCase<"SHAPE_16X32BX2", 4, "shape_16x32bx2">;

def Tcgen05LdStShape: I32EnumAttr<
  "Tcgen05LdStShape",
  "",
  [Tcgen05LdStShape16x64b, Tcgen05LdStShape16x128b, Tcgen05LdStShape16x256b,
   Tcgen05LdStShape32x32b, Tcgen05LdStShape16x32bx2]
> {
  let cppNamespace = "::mlir::NVVM";
  let genSpecializedAttr = 0;
}

def Tcgen05LdStShapeAttr: EnumAttr<NVVM_Dialect, Tcgen05LdStShape, "tcgen05_ldst_shape"> {
  let assemblyFormat = "`<` $value `>`";
}

//===----------------------------------------------------------------------===//
// NVVM tcgen05.ld Op
//===----------------------------------------------------------------------===//

def NVVM_Tcgen05LdOp : NVVM_Op<"tcgen05.ld", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "tensor memory load instructions";
  let arguments = (ins
    // Attributes
    UnitAttr:$pack,
    Tcgen05LdStShapeAttr:$shape,
    // Arguments
    LLVM_PointerTensor:$tmemAddr,
    Optional<I64>:$offset
  );

  let results = (outs AnyTypeOf<[I32, VectorOfLengthAndType<
                                  [2, 4, 8, 16, 32, 64, 128], [I32]>]>:$res);

  let assemblyFormat = [{
    $tmemAddr (`,` $offset^)? (`pack` $pack^)? attr-dict `:` type($res)
  }];

  let description = [{
    Instruction `tcgen05.ld` asynchronously loads data from the Tensor Memory at
    the location specified by the 32-bit address operand `tmemAddr` into the
    destination register `res`, collectively across all threads of the warps.

    The `shape` and the `num` attribute together determines the total
    dimension of the data which is loaded from the Tensor Memory. The `shape`
    attribute indicates the base dimension of data to be accessed as described
    in the Data Movement Shape. The `num` attribute indicates the repeat
    factor on the base dimension resulting in the total dimension of the data
    that is accessed.

    The shape `16x32bx2` performs two accesses into Tensor Memory of the shape
    `16x32b`. The base address of the first access is specified by `tmemAddr`
    and the base address of the second access is specified by
    `tmemAddr + offset`, where `offset` is an immediate argument.

    The unit attribute `pack` can be used to pack two 16-bit
    elements from adjacent columns into a single 32-bit element during the load.

    The following table describes the size of the vector for various combinations
    of `num` and `shape` attributes:
    ```
    |=====================================================================|
    | num/shape      |     16x32bx2/16x64b/32x32b |  16x128b   | 16x256b  |
    |=====================================================================|
    | x1             |          1                 |    2       |    4     |
    | x2             |          2                 |    4       |    8     |
    | x4             |          4                 |    8       |    16    |
    | x8             |          8                 |    16      |    32    |
    | x16            |          16                |    32      |    64    |
    | x32            |          32                |    64      |    128   |
    | x64            |          64                |    128     |    NA    |
    | x128           |          128               |    NA      |    NA    |
    |=====================================================================|
    ```

    Example:
    ```mlir
      nvvm.tcgen05.ld %tmemAddr, %offset pack {
        shape = #nvvm.tcgen05_ldst_shape<shape_16x32bx2>,
      } : <2xi32>
    ```

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st)
  }];

  let hasVerifier = 1;

  string llvmBuilder = [{
    llvm::LLVMContext &Context = moduleTranslation.getLLVMContext();
    auto Pack = llvm::ConstantInt::get(Context, llvm::APInt(1, $pack));

    unsigned num = $_resultType->isVectorTy()
                       ? llvm::cast<llvm::VectorType>($_resultType)
                             ->getElementCount()
                             .getFixedValue()
                       : 1;

    auto ID = getTcgen05LdIntrinsicID($shape, num);
    if (ID == llvm::Intrinsic::not_intrinsic)
      llvm::report_fatal_error("unknow intrinsic signature for tcgen05.ld");

    if ($offset)
      $res = createIntrinsicCall(builder, ID, {$tmemAddr, $offset, Pack});
    else
      $res = createIntrinsicCall(builder, ID, {$tmemAddr, Pack});
  }];
}

//===----------------------------------------------------------------------===//
// NVVM tcgen05.ld.red Op
//===----------------------------------------------------------------------===//

def NVVM_Tcgen05LdRedOp : NVVM_Op<"tcgen05.ld.red",
                          [NVVMRequiresSMaOrSMf<[101, 110], [101, 103, 110]>]> {
  let summary = "Tcgen05 tensor memory load and reduce instructions";
  let arguments = (ins
    Tcgen05LdStShapeAttr:$shape,
    ReductionKindAttr:$op,
    UnitAttr:$abs,
    UnitAttr:$nan,
    LLVM_PointerTensor:$addr,
    Optional<I64>:$offset
  );

  let results = (outs VectorOfLengthAndType<[2, 4, 8, 16, 32, 64, 128],
                                            [I32, F32]>:$data,
                      AnyTypeOf<[I32, F32]>:$redVal);

  let assemblyFormat = [{
    $op $addr (`,` $offset^)? attr-dict `:` type($data) `,` type($redVal)
  }];

  let description = [{
    Instruction `tcgen05.ld.red` asynchronously loads data from the Tensor
    Memory at the location specified by the 32-bit address operand `addr` into
    the destination register `data`, collectively across all threads of the
    warp. The operation also performs reduction operation specified by `op` on
    the loaded data across columns in each lane and stored into `redVal`

    The `shape` and the `num` attribute together determines the total
    dimension of the data which is loaded from the Tensor Memory. The `shape`
    attribute indicates the base dimension of data to be accessed as described
    in the Data Movement Shape. The `num` attribute indicates the repeat
    factor on the base dimension resulting in the total dimension of the data
    that is accessed.

    The shape `16x32bx2` performs two accesses into Tensor Memory of the shape
    `16x32b`. The base address of the first access is specified by `addr`
    and the base address of the second access is specified by
    `addr + offset`, where `offset` is an immediate argument.

    The following table describes the size of the vector for various combinations
    of `num` and `shape` attributes:
    ```
    |=============================================|
    | num/shape      |     16x32bx2/32x32b        |
    |=============================================|
    | x2             |             2              |
    | x4             |             4              |
    | x8             |             8              |
    | x16            |             16             |
    | x32            |             32             |
    | x64            |             64             |
    | x128           |             128            |
    |=============================================|
    ```

    Example:
    ```mlir
      %data, %redval = nvvm.tcgen05.ld.red %addr, %offset {
        shape = #nvvm.tcgen05_ldst_shape<shape_16x32bx2>,
      } : <2xi32>, i32

      %data, %redval = nvvm.tcgen05.ld.red %addr {
        shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>,
      } : <2xf32>, f32
    ```

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld)
  }];

  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];

  string llvmBuilder = [{
    auto [id, args] = NVVM::Tcgen05LdRedOp::getIntrinsicIDAndArgs(
                        *op, moduleTranslation, builder);
    llvm::Value *result = createIntrinsicCall(builder, id, args);
    $data = builder.CreateExtractValue(result, {0});
    $redVal = builder.CreateExtractValue(result, {1});
  }];
}

//===----------------------------------------------------------------------===//
// NVVM tcgen05.st Op
//===----------------------------------------------------------------------===//

def NVVM_Tcgen05StOp : NVVM_Op<"tcgen05.st", [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "tensor memory store instructions";
  let arguments = (ins
    // Attributes
    UnitAttr:$unpack,
    Tcgen05LdStShapeAttr:$shape,
    // Arguments
    LLVM_PointerTensor:$tmemAddr,
    AnyTypeOf<[I32, VectorOfLengthAndType<
                      [2, 4, 8, 16, 32, 64, 128], [I32]>]>:$val,
    Optional<I64>:$offset
  );

  let assemblyFormat = [{
    $tmemAddr `,` $val (`,` $offset^)? (`unpack` $unpack^)? attr-dict `:` type($val)
  }];

  let description = [{
    Instruction `tcgen05.st` asynchronously stores data from the source register `r`
    into the Tensor Memory at the location specified by the 32-bit address operand
    `tmemAddr`, collectively across all threads of the warps.

    The `shape` and the `num` attribute together determines the total dimension of
    the data which is stored to the Tensor Memory. The `shape` indicates the base
    dimension of data to be accessed. The `num` attribute indicates the repeat
    factor on the base dimension resulting in the total dimension of the data that
    is accessed.

    The shape `16x32bx2` performs two accesses into Tensor Memory of the shape
    `16x32b`. The base address of the first access is specified by `tmemAddr`
    and the base address of the second access is specified by
    `tmemAddr + offset`, where `offset` is an immediate argument.

    The unit attribute `unpack` can be used to unpack a 32-bit element
    in the register into two 16-bit elements and store them in adjacent columns.

    The following table describes the size of the vector for various combinations
    of `num` and `shape` attributes:
    ```
    |=====================================================================|
    | num/shape      |     16x32bx2/16x64b/32x32b |  16x128b   | 16x256b  |
    |=====================================================================|
    | x1             |          1                 |    2       |    4     |
    | x2             |          2                 |    4       |    8     |
    | x4             |          4                 |    8       |    16    |
    | x8             |          8                 |    16      |    32    |
    | x16            |          16                |    32      |    64    |
    | x32            |          32                |    64      |    128   |
    | x64            |          64                |    128     |    NA    |
    | x128           |          128               |    NA      |    NA    |
    |=====================================================================|
    ```

    Example:
    ```mlir
      nvvm.tcgen05.st %tmemAddr, %val, %offset unpack {
        shape = #nvvm.tcgen05_ldst_shape<shape_16x32bx2>,
      } : <2xi32>
    ```

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st)
  }];

  string llvmBuilder = [{
    llvm::LLVMContext &Context = moduleTranslation.getLLVMContext();
    auto Unpack = llvm::ConstantInt::get(Context, llvm::APInt(1, $unpack));

    auto valTy = $val->getType();
    uint32_t num = valTy->isVectorTy() ? llvm::cast<llvm::VectorType>(valTy)
                                             ->getElementCount()
                                             .getFixedValue()
                                       : 1;

    auto ID = getTcgen05StIntrinsicID($shape, num);
    if (ID == llvm::Intrinsic::not_intrinsic)
      llvm::report_fatal_error("unknow intrinsic signature for tcgen05.st");

    if ($offset)
      createIntrinsicCall(builder, ID, {$tmemAddr, $offset, $val, Unpack});
    else
      createIntrinsicCall(builder, ID, {$tmemAddr, $val, Unpack});
  }];

  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// NVVM dot.accumulate Ops
//===----------------------------------------------------------------------===//

def DotAccumulateUnsigned : I32EnumAttrCase<"UNSIGNED", 0, "unsigned">;
def DotAccumulateSigned : I32EnumAttrCase<"SIGNED", 1, "signed">;

def DotAccumulateType : I32EnumAttr<"DotAccumulateType",
                              "NVVM DotAccumulateType",
                              [DotAccumulateSigned, DotAccumulateUnsigned]> {
  let cppNamespace = "::mlir::NVVM";
  let genSpecializedAttr = 0;
}

def DotAccumulateTypeAttr : EnumAttr<NVVM_Dialect, DotAccumulateType, "dot_accumulate_type"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_DotAccumulate4WayOp : NVVM_SingleResultIntrinsicOp<"dot.accumulate.4way"> {
  let summary = "Four-way byte dot product-accumulate instruction";
  let description = [{
    Performs a four-way byte dot-product which is accumulated in a 32-bit
    result.
    Operand `a` and `b` are vectors of 4 bytes between which the dot product is 
    computed.

    The `a_type` and `b_type` attributes specify the type of the elements in `a`
    and `b` respectively.
    If `a_type` or `b_type` is `signed`, then the elements in the corresponding 
    vector are sign-extended to 32-bit before the dot product is computed.
    If `a_type` or `b_type` is `unsigned`, then the elements in the 
    corresponding vector are zero-extended to 32-bit instead.

    Operand `c` is a 32-bit integer to which the result is accumulated. It is
    treated as holding a signed integer if any of `a_type` or `b_type` is `s8`.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-dp4a)
  }];
  
  let arguments = (ins
    VectorOfLengthAndType<[4], [I8]>:$a,
    DotAccumulateTypeAttr:$a_type,
    VectorOfLengthAndType<[4], [I8]>:$b,
    DotAccumulateTypeAttr:$b_type,
    I32:$c
  );

  let results = (outs I32:$res);

  let assemblyFormat = "$a $a_type `,` $b $b_type `,` $c attr-dict `:` type($a) `,` type($b)";
}

def NVVM_DotAccumulate2WayOp : NVVM_SingleResultIntrinsicOp<"dot.accumulate.2way"> {
  let summary = "Two-way 16-bit to 8-bit dot product-accumulate instruction";
  let description = [{
    Performs a two-way 16-bit to 8-bit dot-product which is accumulated in a 
    32-bit result.
    Operand `a` is a vector of two 16-bit elements and operand `b` a vector 
    of four 8-bit elements between which the dot product is computed.

    The `a_type` and `b_type` attributes specify the type of the elements in `a`
    and `b` respectively.
    If `a_type` or `b_type` is `s`, then the elements in the corresponding 
    vector are sign-extended to 32-bit before the dot product is computed.
    If `a_type` or `b_type` is `u`, then the elements in the corresponding 
    vector are zero-extended to 32-bit instead.

    The `b_hi` boolean attribute specifies which two bytes of `b` are used for 
    the dot product. If `b_hi` is true, then the dot product is computed 
    between  `a` and elements at indices 2 and 3 of `b`. If `b_hi` is false, 
    then the dot product is computed between `a` and elements at indices 0 and 
    1 of `b`.

    Operand `c` is a 32-bit integer to which the result is accumulated. It is
    treated as holding a signed integer if any of `a_type` or `b_type` is 
    signed.
    
    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#integer-arithmetic-instructions-dp2a)
  }];

  let arguments = (ins
    VectorOfLengthAndType<[2], [I16]>:$a,
    DotAccumulateTypeAttr:$a_type,
    VectorOfLengthAndType<[4], [I8]>:$b,
    DotAccumulateTypeAttr:$b_type,
    I32:$c,
    BoolAttr:$b_hi
  );

  let results = (outs I32:$res);

  let assemblyFormat = "$a $a_type `,` $b $b_type `,` $c attr-dict `:` type($a) `,` type($b)";
}

//===----------------------------------------------------------------------===//
// NVVM clusterlaunchcontrol Ops.
//===----------------------------------------------------------------------===//

def NVVM_ClusterLaunchControlTryCancelOp
    : NVVM_Op<"clusterlaunchcontrol.try.cancel", [NVVMRequiresSM<100>]> {
  let summary = "Request atomically canceling the launch of a cluster that has not started running yet";
  let description = [{
    `clusterlaunchcontrol.try.cancel` requests atomically canceling the launch 
    of a cluster that has not started running yet. It asynchronously writes an 
    opaque response to shared memory indicating whether the operation succeeded 
    or failed.

    Operand `smemAddress` specifies the naturally aligned address of the 
    16-byte wide shared memory location where the request's response is written.

    Operand `mbarrier` specifies the mbarrier object used to track the 
    completion of the asynchronous operation.

    If `multicast` is specified, the response is asynchronously written to the 
    corresponding local shared memory location (specifed by `addr`) of each CTA 
    in the requesting cluster.

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-try-cancel)
  }];

  let arguments = (ins UnitAttr:$multicast, 
                       LLVM_PointerShared: $smemAddress,
                       LLVM_PointerShared: $mbarrier);

  let assemblyFormat = "(`multicast` $multicast^ `,`)? $smemAddress `,` $mbarrier attr-dict";

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];
  
  string llvmBuilder = [{
    auto [id, args] = 
    NVVM::ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs(
                        *op, moduleTranslation, builder);
    createIntrinsicCall(builder, id, args);
  }];
}

def ClusterLaunchControlIsCanceled
  : I32EnumCase<"IS_CANCELED", 0, "is_canceled">;
def ClusterLaunchControlGetFirstCTAIDX
  : I32EnumCase<"GET_FIRST_CTA_ID_X", 1, "get_first_cta_id_x">;
def ClusterLaunchControlGetFirstCTAIDY
  : I32EnumCase<"GET_FIRST_CTA_ID_Y", 2, "get_first_cta_id_y">;
def ClusterLaunchControlGetFirstCTAIDZ
  : I32EnumCase<"GET_FIRST_CTA_ID_Z", 3, "get_first_cta_id_z">;

def ClusterLaunchControlQueryType
  : I32Enum<"ClusterLaunchControlQueryType",
      "NVVM ClusterLaunchControlQueryType", 
      [ClusterLaunchControlIsCanceled, ClusterLaunchControlGetFirstCTAIDX,
      ClusterLaunchControlGetFirstCTAIDY, ClusterLaunchControlGetFirstCTAIDZ]> {
  let cppNamespace = "::mlir::NVVM";
}

def ClusterLaunchControlQueryTypeAttr
  : EnumAttr<NVVM_Dialect,
      ClusterLaunchControlQueryType, "cluster_launch_control_query_type"> {
  let assemblyFormat = "$value";
}

def NVVM_ClusterLaunchControlQueryCancelOp
    : NVVM_Op<"clusterlaunchcontrol.query.cancel",
              [NVVMRequiresSM<100>, InferTypeOpAdaptor]> {
  let summary = "Query the response of a clusterlaunchcontrol.try.cancel operation";
  let description = [{
    `clusterlaunchcontrol.query.cancel` queries the response of a 
    `clusterlaunchcontrol.try.cancel` operation specified by operand 
    `try_cancel_response`.

    Operand `query_type` specifies the type of query to perform and can be one 
    of the following:
    - `is_canceled` : Returns true if the try cancel request succeeded, 
    and false otherwise.
    - `get_first_cta_id_{x/y/z}` : Returns the x, y, or z coordinate of the 
    first CTA in the canceled cluster. Behaviour is defined only if the try 
    cancel request succeeded. 

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-query-cancel)
  }];

  let arguments = (ins ClusterLaunchControlQueryTypeAttr:$query_type, 
                       I128:$try_cancel_response);
  let results = (outs AnyTypeOf<[I1, I32]>:$res);
                                 
  let assemblyFormat = "`query` `=` $query_type `,` $try_cancel_response attr-dict `:` type($res)";
  
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];
  
  string llvmBuilder = [{
    auto [id, args] = 
    NVVM::ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs(
                        *op, moduleTranslation, builder);
    $res = createIntrinsicCall(builder, id, args);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM tcgen05.mma Ops
//===----------------------------------------------------------------------===//

def Tcgen05MMAKindF16          : I32EnumAttrCase<"F16",      0, "f16">;
def Tcgen05MMAKindTF32         : I32EnumAttrCase<"TF32",     1, "tf32">;
def Tcgen05MMAKindF8F6F4       : I32EnumAttrCase<"F8F6F4",   2, "f8f6f4">;
def Tcgen05MMAKindI8           : I32EnumAttrCase<"I8",       3, "i8">;
def Tcgen05MMAKindMXF8F6F4     : I32EnumAttrCase<"MXF8F6F4", 4, "mxf8f6f4">;
def Tcgen05MMAKindMXF4         : I32EnumAttrCase<"MXF4",     5, "mxf4">;
def Tcgen05MMAKindMXF4NVF4     : I32EnumAttrCase<"MXF4NVF4", 6, "mxf4nvf4">;

def Tcgen05MMAKind : I32EnumAttr<
  "Tcgen05MMAKind",
  "tcgen05 MMA Supported Types",
  [Tcgen05MMAKindF16, Tcgen05MMAKindTF32, Tcgen05MMAKindF8F6F4,
   Tcgen05MMAKindI8, Tcgen05MMAKindMXF8F6F4, Tcgen05MMAKindMXF4,
   Tcgen05MMAKindMXF4NVF4]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}

def Tcgen05MMAKindAttr : EnumAttr<NVVM_Dialect, Tcgen05MMAKind, "tcgen05_mma_kind"> {
  let description = [{
    The Tcgen05MMAKind attribute describes the allowed set of types for matrix A and B in the tcgen05.mma Ops.
    The following are supported types for each kind:

    ```
    +-------------+--------------------------------------------+
    | Matrix Kind |      supported types for A / B             |
    +-------------+--------------------------------------------+
    | f16         | f16, bf16                                  |
    | tf32        | tf32                                       |
    | f8f6f4      | e4m3, e5m2, e2m3, e3m2, e2m1               |
    | i8          | unsigned 8b, signed 8b                     |
    | mxf8f6f4    | e4m3, e5m2, e2m3, e3m2, e2m1               |
    | mxf4        | e2m1                                       |
    | mxf4nvf4    | e2m1                                       |
    +-------------+--------------------------------------------+
    ```
  }];
  let assemblyFormat = "`<` $value `>`";
}

def Tcgen05MMACollectorOpDiscard  : I32EnumAttrCase<"DISCARD", 0, "discard">;
def Tcgen05MMACollectorOpLastUse  : I32EnumAttrCase<"LASTUSE", 1, "lastuse">;
def Tcgen05MMACollectorOpFill     : I32EnumAttrCase<"FILL",    2, "fill">;
def Tcgen05MMACollectorOpUse      : I32EnumAttrCase<"USE",     3, "use">;

def Tcgen05MMACollectorOp : I32EnumAttr<
  "Tcgen05MMACollectorOp",
  "tcgen05.mma Collector Buffer Operation",
  [Tcgen05MMACollectorOpDiscard,
   Tcgen05MMACollectorOpLastUse,
   Tcgen05MMACollectorOpFill,
   Tcgen05MMACollectorOpUse]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}

def Tcgen05MMACollectorOpAttr : EnumAttr<NVVM_Dialect, Tcgen05MMACollectorOp, "tcgen05_mma_collectorop"> {
  let description = [{
    Tcgen05MMACollectorOp attribute specifies the collector buffer operations.
    The following are the supported operations:
      * discard : Release buffer after use (default)
      * lastuse : Mark buffer for last use
      * fill    : Fill buffer
      * use     : Use buffer without modification
  }];
  let assemblyFormat = "`<` $value `>`";
}

defvar Tcgen05MMANonBlockScaleKindList =
  [Tcgen05MMAKindF16, Tcgen05MMAKindTF32,
   Tcgen05MMAKindF8F6F4, Tcgen05MMAKindI8];

defvar Tcgen05MMABlockScaleKindList =
  [Tcgen05MMAKindMXF8F6F4, Tcgen05MMAKindMXF4, Tcgen05MMAKindMXF4NVF4];

defvar Tcgen05MMANonBlockScaleKindAttr =
  ConfinedAttr<Tcgen05MMAKindAttr,
    [EnumAttrIsOneOf<Tcgen05MMAKindAttr, Tcgen05MMANonBlockScaleKindList>]>;

defvar Tcgen05MMABlockScaleKindAttr =
  ConfinedAttr<Tcgen05MMAKindAttr,
    [EnumAttrIsOneOf<Tcgen05MMAKindAttr, Tcgen05MMABlockScaleKindList>]>;

def NVVM_Tcgen05MMAOp : NVVM_Op<"tcgen05.mma",
    [AttrSizedOperandSegments, NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs MMA operation on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma` operation is an asynchronous tensor core instruction that
    performs matrix multiplication, accumulation in a single fused operation. It
    targets 5th-generation tensor cores, providing developers with fine-grained
    control over execution and scheduling.

    ```
    D = A * B + (D * 2^ -scaleInputD)    // if `scaleInputD` is provided
    D = A * B                            // if `enableInputD` is false
    D = A * B + D                        // otherwise
    ```

    where:
    - A is an `M x K` matrix in tensor memory or described using shared memory descriptor
    - B is a `K x N` matrix described using shared memory descriptor
    - D is an `M x N` accumulator matrix in tensor memory

    The `shared memory descriptor` can be generated using `tcgen05.mma_smem_desc` Op

    - idesc is a 32-bit value representing the [Instruction Descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instruction-descriptor)

    Optional Operands:
    - `scaleInputD` is an Immediate value operand used for scaling D matrix by 2 ^ (-scaleInputD). The valid range is [0, 15]

    - `disableOutputLane` is a vector mask for selective output
      * vector<4 x i32> when ctaGroup is CTA_1
      * vector<8 x i32> when ctaGroup is CTA_2

    Required Attributes:
    - `kind` is a Tcgen05MMAKind attribute

    - `ctaGroup` specifies CTA group configuration
      * cta_1: MMA will be performed on the current thread's CTA
      * cta_2: MMA will be performed on the current thread and it's peer CTA

    Default Attributes:
    - collectorOp is a Tcgen05MMACollectorOp attribute with matrix A as the collector buffer

    - `aShift` shifts the rows of the A matrix down by one row and can only be
       applied if A is in tensor memory

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma)
  }];

  let arguments = (ins
      Tcgen05MMANonBlockScaleKindAttr:$kind,
      CTAGroupKindAttr:$ctaGroup,
      DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                        "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
      UnitAttr:$aShift,
      LLVM_PointerTensor:$matrixD,
      AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
      I64:$matrixB,
      I32:$idesc,
      I1:$enableInputD,
      Optional<I64>:$scaleInputD,
      Optional<FixedVectorOfLengthAndType<[4, 8], [I32]>>:$disableOutputLane
    );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD (`scale` `=` $scaleInputD^)?
    (`mask` `=` $disableOutputLane^)? attr-dict `:` `(` type(operands) `)`
  }];

  let hasVerifier = true;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair getIntrinsicIDAndArgs(
        Operation &op, LLVM::ModuleTranslation &mt,
        llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
    auto [ID, args] = NVVM::Tcgen05MMAOp::getIntrinsicIDAndArgs(
        *op, moduleTranslation, builder);
    createIntrinsicCall(builder, ID, args);
  }];
}

def NVVM_Tcgen05MMASparseOp : NVVM_Op<"tcgen05.mma.sp",
    [AttrSizedOperandSegments, NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs MMA operation with sparse A matrix on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma.sp` operation is an asynchronous tensor core instruction
    that performs matrix multiplication, accumulation with sparse `A` matrix in
    a single fused operation. It targets 5th-generation tensor cores, providing
    developers with fine-grained control over execution and scheduling.

    ```
    D = A * B + (D * 2^ -scaleInputD)    // if `scaleInputD` is provided
    D = A * B                            // if `enableInputD` is false
    D = A * B + D                        // otherwise
    ```

    where:
    - A is an `M x (K / 2)` matrix in tensor memory or described using shared memory descriptor
    - B is a `K x N` matrix described using shared memory descriptor
    - D is an `M x N` accumulator matrix in tensor memory
    - sparseMetadata located in tensor memory specifies the mapping of the `K / 2`
    non-zero elements to the K elements before performing the MMA operation

    Other attributes and operands are similar to that of tcgen05.mma Op

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-sp)
  }];

  let arguments = (ins
    Tcgen05MMANonBlockScaleKindAttr:$kind,
    CTAGroupKindAttr:$ctaGroup,
    DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                      "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
    UnitAttr:$aShift,
    LLVM_PointerTensor:$matrixD,
    AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
    I64:$matrixB,
    I32:$idesc,
    I1:$enableInputD,
    LLVM_PointerTensor:$sparseMetadata,
    Optional<I64>:$scaleInputD,
    Optional<FixedVectorOfLengthAndType<[4, 8], [I32]>>:$disableOutputLane
  );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD `,` $sparseMetadata (`scale` `=` $scaleInputD^)? (`mask` `=` $disableOutputLane^)? attr-dict `:` `(` type(operands) `)`
  }];

  let hasVerifier = true;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair getIntrinsicIDAndArgs(
        Operation &op, LLVM::ModuleTranslation &mt,
        llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
    auto [ID, args] = NVVM::Tcgen05MMASparseOp::getIntrinsicIDAndArgs(
        *op, moduleTranslation, builder);
    createIntrinsicCall(builder, ID, args);
  }];
}

def Tcgen05MMABlockScaleDefault : I32EnumAttrCase<"DEFAULT", 0, "default">;
def Tcgen05MMABlockScaleBlock16      : I32EnumAttrCase<"BLOCK16", 1, "block16">;
def Tcgen05MMABlockScaleBlock32      : I32EnumAttrCase<"BLOCK32", 2, "block32">;

def Tcgen05MMABlockScale
    : I32EnumAttr<"Tcgen05MMABlockScale",
                  "tcgen05.mma block scale attribute",
                  [Tcgen05MMABlockScaleDefault, Tcgen05MMABlockScaleBlock16,
                   Tcgen05MMABlockScaleBlock32]> {
  let cppNamespace = "::mlir::NVVM";
  let genSpecializedAttr = 0;
}

def Tcgen05MMABlockScaleAttr : EnumAttr<NVVM_Dialect, Tcgen05MMABlockScale,
                                          "tcgen05_mma_block_scale"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_Tcgen05MMABlockScaleOp : NVVM_Op<"tcgen05.mma.block_scale",
    [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs block scaled MMA operation on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma.block_scale` operation is an asynchronous tensor core instruction
    that performs matrix multiplication, accumulation with block scaling in a
    single fused operation. It targets 5th-generation tensor cores, providing
    developers with fine-grained control over execution and scheduling.

    ```
    D = (A * scale_a)  * (B * scale_b)`      // if `enableInputD` is false
    D = (A * scale_a)  * (B * scale_b) + D`
    ```

    where:
    - A is an M x (K / 2) matrix in tensor memory or described using shared memory descriptor
    - B is a K x N matrix described using shared memory descriptor
    - D is an M x N accumulator matrix in tensor memory
    - `scale_a` and `scale_b` are matrices in tensor memory used to scale `A` and `B` respectively

    The `shared memory descriptor` can be generated using `tcgen05.mma_smem_desc` Op

    - `idesc` is a 32 bit value representing the [Instruction Descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instruction-descriptor)

    Required Attributes:
    - `kind` is a Tcgen05MMAKind attribute restricted to mxf8f6f4, mxf4, or mxf4nvf4

    - `ctaGroup` specifies CTA group configuration
      * cta_1: MMA will be performed on the current thread's CTA
      * cta_2: MMA will be performed on the current thread and it's peer CTA

    Default Attributes:
    - collectorOp is a Tcgen05MMACollectorOp attribute with matrix A as the collector buffer

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma)
  }];

  let arguments = (ins
      Tcgen05MMABlockScaleKindAttr:$kind,
      CTAGroupKindAttr:$ctaGroup,
      DefaultValuedAttr<Tcgen05MMABlockScaleAttr,
                      "Tcgen05MMABlockScale::DEFAULT">:$blockScale,
      DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                        "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
      LLVM_PointerTensor:$matrixD,
      AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
      I64:$matrixB,
      I32:$idesc, I1:$enableInputD,
      LLVM_PointerTensor:$scaleA,
      LLVM_PointerTensor:$scaleB
    );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD `,`
    $scaleA `,` $scaleB attr-dict `:` `(` type(operands) `)`
  }];

  let hasVerifier = true;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
    auto [ID, args] = NVVM::Tcgen05MMABlockScaleOp::getIntrinsicIDAndArgs(
        *op, moduleTranslation, builder);
    createIntrinsicCall(builder, ID, args);
  }];
}

def NVVM_Tcgen05MMASparseBlockScaleOp : NVVM_Op<"tcgen05.mma.sp.block_scale",
    [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs block scaled MMA operation with sparse A matrix on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma.sp.block_scale` operation is an asynchronous tensor core
    instruction that performs matrix multiplication, accumulation with block
    scaling, and sparse `A` matrix in a single fused operation. It targets
    5th-generation tensor cores, providing developers with fine-grained control
    over execution, and scheduling.

    ```
    D = (A * scale_a)  * (B * scale_b)      // if `enableInputD` is specified
    D = (A * scale_a)  * (B * scale_b) + D  // otherwise
    ```

    where:
    - A is an M x (K / 2) matrix in tensor memory or described using shared memory descriptor
    - B is a K x N matrix described using shared memory descriptor
    - D is an M x N accumulator matrix in tensor memory
    - `scale_a` and `scale_b` are matrices in tensor memory used to scale `A` and `B` respectively

    Other attributes and operands are similar to that of tcgen05.mma.block_scale Op

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-sp)
  }];

  let arguments = (ins
    Tcgen05MMABlockScaleKindAttr:$kind,
    CTAGroupKindAttr:$ctaGroup,
    DefaultValuedAttr<Tcgen05MMABlockScaleAttr,
                      "Tcgen05MMABlockScale::DEFAULT">:$blockScale,
    DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                      "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
    LLVM_PointerTensor:$matrixD,
    AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
    I64:$matrixB,
    I32:$idesc,
    I1:$enableInputD,
    LLVM_PointerTensor:$sparseMetadata,
    LLVM_PointerTensor:$scaleA,
    LLVM_PointerTensor:$scaleB
  );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD `,`
    $sparseMetadata `,` $scaleA `,` $scaleB attr-dict `:` `(` type(operands) `)`
  }];

  let hasVerifier = true;

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
    auto [ID, args] = NVVM::Tcgen05MMASparseBlockScaleOp::getIntrinsicIDAndArgs(
        *op, moduleTranslation, builder);
    createIntrinsicCall(builder, ID, args);
  }];
}

def Tcgen05MMACollectorBBuffer0  : I32EnumAttrCase<"B0", 0, "b0">;
def Tcgen05MMACollectorBBuffer1  : I32EnumAttrCase<"B1", 1, "b1">;
def Tcgen05MMACollectorBBuffer2  : I32EnumAttrCase<"B2", 2, "b2">;
def Tcgen05MMACollectorBBuffer3  : I32EnumAttrCase<"B3", 3, "b3">;

def Tcgen05MMACollectorBBuffer : I32EnumAttr<
  "Tcgen05MMACollectorBBuffer",
  "tcgen05 MMA Collector Buffer B Attribute",
  [Tcgen05MMACollectorBBuffer0, Tcgen05MMACollectorBBuffer1, Tcgen05MMACollectorBBuffer2,
   Tcgen05MMACollectorBBuffer3]> {
    let cppNamespace = "::mlir::NVVM";
    let genSpecializedAttr = 0;
}

def Tcgen05MMACollectorBBufferAttr : EnumAttr<NVVM_Dialect, Tcgen05MMACollectorBBuffer, "tcgen05_mma_collectorb"> {
  let assemblyFormat = "`<` $value `>`";
}

def NVVM_Tcgen05MMAWsOp : NVVM_Op<"tcgen05.mma.ws",
    [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs weight stationary convolution MMA operation on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma.ws` operation is an asynchronous tensor core instruction
    that performs weight stationary convolution matrix multiplication, accumulation
    in a single fused operation. It targets 5th-generation tensor cores, providing
    developers with fine-grained control over execution, and scheduling.

    ```
    D = A * B`      // if `enableInputD` is false
    D = A * B + D`  // otherwise
    ```

    where:
    - A is an `M x K` matrix in tensor memory or described using shared memory descriptor
    - B is a `K x N` matrix described using shared memory descriptor
    - D is an `M x N` accumulator matrix in tensor memory

    The `shared memory descriptor` can be generated using `tcgen05.mma_smem_desc` Op

    - idesc is a 32-bit value representing the [Instruction Descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instruction-descriptor)

    Optional Operands:
    - zeroColMask is a 64 bit value representing the [Zero-column mask descriptor](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-zero-column-mask-descriptor)

    Required Attributes:
    - `kind` is a Tcgen05MMAKind attribute

    Default Valued Attributes:
    - collectorBBuffer specifies collector buffer for matrix B: b0 (default), b1, b2, b3

    - collectorOp is a Tcgen05MMACollectorOp attribute with matrix B as the collector buffer

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-ws)
  }];

  let arguments = (ins
    Tcgen05MMANonBlockScaleKindAttr:$kind,
    DefaultValuedAttr<Tcgen05MMACollectorBBufferAttr,
                      "Tcgen05MMACollectorBBuffer::B0">:$collectorBBuffer,
    DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                      "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
    LLVM_PointerTensor:$matrixD,
    AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
    I64:$matrixB,
    I32:$idesc,
    I1:$enableInputD,
    Optional<I64>:$zeroColMask
  );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD (`,` $zeroColMask^)?
    attr-dict `:` `(` type(operands) `)`
  }];

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair getIntrinsicIDAndArgs(
        Operation &op, LLVM::ModuleTranslation &mt,
        llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
    auto [ID, args] =
        NVVM::Tcgen05MMAWsOp::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
    createIntrinsicCall(builder, ID, args);
  }];
}

def NVVM_Tcgen05MMAWsSparseOp : NVVM_Op<"tcgen05.mma.ws.sp",
    [NVVMRequiresSMf<[100, 101, 110]>]> {
  let summary = "Performs weight stationary convolution MMA with sparse A matrix on 5th-gen tensor cores";

  let description = [{
    The `tcgen05.mma.ws.sp` operation is an asynchronous tensor core instruction
    that performs weight stationary convolution matrix multiplication, accumulation
    with sparse `A` matrix in a single fused operation. It targets 5th-generation
    tensor cores, providing developers with fine-grained control over execution,
    and scheduling.

    ```
    D = A * B`      // if `enableInputD` is false
    D = A * B + D`  // otherwise
    ```

    where:
    - A is an M x (K / 2) matrix in memory or descriptor format
    - B is a K x N matrix
    - D is an M x N accumulator matrix
    - sparseMetadata located in tensor memory specifies the mapping of the `K / 2`
    non-zero elements to the K elements before performing the MMA operation

    Other attributes and operands are similar to that of tcgen05.mma.ws Op

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-ws-sp)
  }];

  let arguments = (ins
    Tcgen05MMANonBlockScaleKindAttr:$kind,
    DefaultValuedAttr<Tcgen05MMACollectorBBufferAttr,
                      "Tcgen05MMACollectorBBuffer::B0">:$collectorBBuffer,
    DefaultValuedAttr<Tcgen05MMACollectorOpAttr,
                      "Tcgen05MMACollectorOp::DISCARD">:$collectorOp,
    LLVM_PointerTensor:$matrixD,
    AnyTypeOf<[LLVM_PointerTensor, I64]>:$matrixA,
    I64:$matrixB,
    I32:$idesc,
    I1:$enableInputD,
    LLVM_PointerTensor:$sparseMetadata,
    Optional<I64>:$zeroColMask
  );

  let assemblyFormat = [{
    $matrixD `,` $matrixA `,` $matrixB `,` $idesc `,` $enableInputD `,` $sparseMetadata (`,` $zeroColMask^)? attr-dict `:` `(` type(operands) `)`
  }];

  let extraClassDeclaration = [{
    static mlir::NVVM::IDArgPair
    getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                          llvm::IRBuilderBase &builder);
  }];

  let llvmBuilder = [{
      auto [ID, args] = NVVM::Tcgen05MMAWsSparseOp::getIntrinsicIDAndArgs(
          *op, moduleTranslation, builder);
      createIntrinsicCall(builder, ID, args);
  }];
}

def SIMTFloatType : AnyTypeOf<[F16, BF16, F32, F64,
                      VectorOfLengthAndType<[2], [F16, BF16, F32, F64]>]>;

def SaturationModeSatOrNone : 
  ConfinedAttr<SaturationModeAttr, [EnumAttrIsOneOf<SaturationModeAttr, 
                [SaturationModeNone, SaturationModeSat]>]>;

def FPArithRoundingMode : 
  ConfinedAttr<FPRoundingModeAttr, [EnumAttrIsOneOf<FPRoundingModeAttr, 
                [FPRoundingModeNone, FPRoundingModeRM, FPRoundingModeRN, 
                 FPRoundingModeRP, FPRoundingModeRZ]>]>;

class NVVM_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
    NVVM_Op<mnemonic, traits # [Pure, SameOperandsAndResultType]>,
    Arguments<(ins SIMTFloatType:$lhs, SIMTFloatType:$rhs,
      DefaultValuedAttr<FPArithRoundingMode, "FPRoundingMode::NONE">:$rnd,
      DefaultValuedAttr<SaturationModeSatOrNone, "SaturationMode::NONE">:$sat,
      DefaultValuedAttr<BoolAttr, "false">:$ftz)>,
    Results<(outs SIMTFloatType:$res)> {
  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($res)";
}

def NVVM_AddFOp : NVVM_FloatBinaryOp<"addf", [Commutative]> {
  let summary = [{
    Performs floating point addition of the given arguments `lhs` and `rhs`
  }];
  let description = [{
    The `nvvm.addf` operation performs floating point addition of two floating 
    point operands of the same type.

    The rounding mode is specified by the `rnd` attribute, saturation mode by 
    the `sat` attribute, and flush-to-zero by the `ftz` attribute.

    For more information, see PTX ISA:
    - [floating point addition](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-add)
    - [half-precision floating point addition](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-add)
  }];
  let hasVerifier = 1;
  
  let extraClassDeclaration = [{
    // Shared lowering for nvvm.addf and nvvm.subf (via fneg+add)
    static void lowerAddFToLLVMIR(
        llvm::Value *lhs, llvm::Value *rhs, Value res,
        NVVM::FPRoundingMode rnd, NVVM::SaturationMode sat, bool ftz,
        LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];
  let llvmBuilder = [{
    NVVM::AddFOp::lowerAddFToLLVMIR($lhs, $rhs, op.getRes(), $rnd, $sat, $ftz, 
                                    moduleTranslation, builder);
  }];
}

def NVVM_SubFOp : NVVM_FloatBinaryOp<"subf"> {
  let summary = [{
    Performs floating point subtraction of the given arguments `lhs` and `rhs`
  }];
  let description = [{
    The `nvvm.subf` operation performs floating point subtraction of two 
    operands.

    It supports the same type combinations and modifiers as `nvvm.addf`.
    This is equivalent to `nvvm.addf(lhs, -rhs)`.
    
    For more information, see PTX ISA:
    - [floating point subtraction](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sub) 
    - [half-precision floating point subtraction](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-sub)
 
  }];
  let hasCanonicalizer = 1;
  let hasVerifier = 1;

  let llvmBuilder = [{
    // sub(a, b) = add(a, -b)
    llvm::Value *rhs = builder.CreateFNeg($rhs);
    NVVM::AddFOp::lowerAddFToLLVMIR($lhs, rhs, op.getRes(), $rnd, $sat, $ftz, 
                                    moduleTranslation, builder);
  }];
}

def NVVM_FmaOp : NVVM_Op<"fma", [Pure, SameOperandsAndResultType]> {
  let summary = [{
    Performs floating point fused multiply-add operation with support for mixed 
    precision operands
  }];
  let description = [{
    The `nvvm.fma` operation performs floating point fused multiply-add of 
    three operands of the same type.

    The rounding mode is specified by the `rnd` attribute, saturation mode by 
    the `sat` attribute, flush-to-zero by the `ftz` attribute, and ReLU by the 
    `relu` attribute.
    
    Out-of-bounds (OOB) behavior is controlled by the `oob` attribute. `oob` 
    clamps the result to 0 if either of the operands is `OOB NaN` (see [Tensors](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensors)).

    For more information, see PTX ISA:
    - [floating point fused multiply-add](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma)
    - [half-precision floating point fused multiply-add](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma)
  }];
  let arguments = (ins
    SIMTFloatType:$a,
    SIMTFloatType:$b,
    SIMTFloatType:$c,
    FPArithRoundingMode:$rnd,
    DefaultValuedAttr<SaturationModeSatOrNone, "SaturationMode::NONE">:$sat,
    DefaultValuedAttr<BoolAttr, "false">:$ftz,
    DefaultValuedAttr<BoolAttr, "false">:$relu,
    DefaultValuedAttr<BoolAttr, "false">:$oob
  );
  let results = (outs SIMTFloatType:$res);
  let assemblyFormat = "$a `,` $b `,` $c attr-dict `:` type($a)";
  let hasVerifier = 1;

  let extraClassDeclaration = [{
    static void lowerFmaToLLVMIR(
        Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
  }];
  let llvmBuilder = [{
    NVVM::FmaOp::lowerFmaToLLVMIR(*op, moduleTranslation, builder);
  }];
}

//===----------------------------------------------------------------------===//
// NVVM sqrt op definitions
//===----------------------------------------------------------------------===//

def NVVM_SqrtOp
    : NVVM_SingleResultIntrinsicOp<"sqrt", [Pure, SameOperandsAndResultType]> {
  let summary = "Take the square root of a value";
  let description = [{
    Compute sqrt(src) and store the result in res.

    For more information, see PTX ISA:
    [sqrt](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt)
  }];
  let arguments = (ins AnyTypeOf<[F32, F64]>:$src, FPArithRoundingMode:$rnd,
      DefaultValuedAttr<BoolAttr, "false">:$ftz);
  let results = (outs AnyTypeOf<[F32, F64]>:$res);
  let assemblyFormat = "$src attr-dict `:` type($src)";
  let hasVerifier = 1;
}

def NVVM_SqrtApproxOp : NVVM_F32UnaryApproxOp<"sqrt.approx"> {
  let summary = "Square root (fast approximation)";
  let description = [{
    Computes a fast approximation of the square root of the input value
    (`res = sqrt(src)`). The maximum relative error over the entire positive
    finite range is 2^-23.

    The `ftz` attribute, when set, flushes subnormal inputs and results to
    sign-preserving zero.

    For more information, see PTX ISA:
    [sqrt](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt)
  }];
}

//===----------------------------------------------------------------------===//
// NVVM div op definitions
//===----------------------------------------------------------------------===//

def NVVM_DivFOp
    : NVVM_SingleResultIntrinsicOp<"divf", [Pure, SameOperandsAndResultType]> {
  let summary = "Divide one value by another";
  let description = [{
    Divides lhs by rhs, stores result in res (`res = lhs / rhs`).

    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-div)
  }];
  let arguments = (ins AnyTypeOf<[F32, F64]>:$lhs, AnyTypeOf<[F32, F64]>:$rhs,
      DefaultValuedAttr<FPArithRoundingMode, "FPRoundingMode::NONE">:$rnd,
      DefaultValuedAttr<BoolAttr, "false">:$ftz,
      DefaultValuedAttr<BoolAttr, "false">:$approx,
      DefaultValuedAttr<BoolAttr, "false">:$full);
  let results = (outs AnyTypeOf<[F32, F64]>:$res);
  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($res)";
  let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
// NVVM tensormap.replace Op
//===----------------------------------------------------------------------===//

// Tensormap Field Kind Enum
def TensormapGlobalAddressField : I32EnumCase<"GLOBAL_ADDRESS", 0, "global_address">;
def TensormapRankField : I32EnumCase<"RANK", 1, "rank">;
def TensormapBoxDimField : I32EnumCase<"BOX_DIM", 2, "box_dim">;
def TensormapGlobalDimField : I32EnumCase<"GLOBAL_DIM", 3, "global_dim">;
def TensormapGlobalStrideField : I32EnumCase<"GLOBAL_STRIDE", 4, "global_stride">;
def TensormapElementStrideField : I32EnumCase<"ELEMENT_STRIDE", 5, "element_stride">;
def TensormapElemtypeField : I32EnumCase<"ELEMTYPE", 6, "elemtype">;
def TensormapInterleaveLayoutField : I32EnumCase<"INTERLEAVE_LAYOUT", 7, "interleave_layout">;
def TensormapSwizzleModeField : I32EnumCase<"SWIZZLE_MODE", 8, "swizzle_mode">;
def TensormapSwizzleAtomicityField : I32EnumCase<"SWIZZLE_ATOMICITY", 9, "swizzle_atomicity">;
def TensormapFillModeField : I32EnumCase<"FILL_MODE", 10, "fill_mode">;

def TensormapField : I32Enum<"TensormapField", "NVVM Tensormap Field Kind",
  [TensormapGlobalAddressField, TensormapRankField, TensormapBoxDimField, 
   TensormapGlobalDimField, TensormapGlobalStrideField, 
   TensormapElementStrideField, TensormapElemtypeField, 
   TensormapInterleaveLayoutField, TensormapSwizzleModeField, 
   TensormapSwizzleAtomicityField, TensormapFillModeField]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapFieldAttr : EnumAttr<NVVM_Dialect, TensormapField, "tensormap_field"> {
  let assemblyFormat = "$value";
}

// Tensormap Elemtype Enum
def TensormapElemtypeU8 : I32EnumCase<"U8", 0, "u8">;
def TensormapElemtypeU16 : I32EnumCase<"U16", 1, "u16">;
def TensormapElemtypeU32 : I32EnumCase<"U32", 2, "u32">;
def TensormapElemtypeS32 : I32EnumCase<"S32", 3, "s32">;
def TensormapElemtypeU64 : I32EnumCase<"U64", 4, "u64">;
def TensormapElemtypeS64 : I32EnumCase<"S64", 5, "s64">;
def TensormapElemtypeF16 : I32EnumCase<"F16", 6, "f16">;
def TensormapElemtypeF32 : I32EnumCase<"F32", 7, "f32">;
def TensormapElemtypeF32_Ftz : I32EnumCase<"F32_FTZ", 8, "f32.ftz">;
def TensormapElemtypeF64 : I32EnumCase<"F64", 9, "f64">;
def TensormapElemtypeBF16 : I32EnumCase<"BF16", 10, "bf16">;
def TensormapElemtypeTF32 : I32EnumCase<"TF32", 11, "tf32">;
def TensormapElemtypeTF32_Ftz : I32EnumCase<"TF32_FTZ", 12, "tf32.ftz">;
def TensormapElemtypeB4x16 : I32EnumCase<"B4x16", 13, "b4x16">;
def TensormapElemtypeB4x16P64 : I32EnumCase<"B4x16P64", 14, "b4x16_p64">;
def TensormapElemtypeB6x16P32 : I32EnumCase<"B6x16P32", 15, "b6x16_p32">;

def TensormapElemtype : I32Enum<"TensormapElemtype", "NVVM Tensormap Elemtype",
  [TensormapElemtypeU8, TensormapElemtypeU16, 
   TensormapElemtypeU32,TensormapElemtypeS32, 
   TensormapElemtypeU64, TensormapElemtypeS64,
   TensormapElemtypeF16, TensormapElemtypeF32, 
   TensormapElemtypeF32_Ftz,TensormapElemtypeF64, 
   TensormapElemtypeBF16, TensormapElemtypeTF32,
   TensormapElemtypeTF32_Ftz, TensormapElemtypeB4x16, 
   TensormapElemtypeB4x16P64, TensormapElemtypeB6x16P32]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapElemtypeAttr : EnumAttr<NVVM_Dialect, TensormapElemtype, "tensormap_elemtype"> {
  let assemblyFormat = "`<` $value `>`";
}

// Tensormap Interleave Layout Enum
def TensormapInterleaveLayoutNoInterleave : 
  I32EnumCase<"NO_INTERLEAVE", 0, "no_interleave">;
def TensormapInterleaveLayoutB16 : I32EnumCase<"B16", 1, "b16">;
def TensormapInterleaveLayoutB32 : I32EnumCase<"B32", 2, "b32">;

def TensormapInterleaveLayout : 
  I32Enum<"TensormapInterleaveLayout", "NVVM Tensormap Interleave Layout",
  [TensormapInterleaveLayoutNoInterleave, 
   TensormapInterleaveLayoutB16,
   TensormapInterleaveLayoutB32]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapInterleaveLayoutAttr : EnumAttr<NVVM_Dialect,   
  TensormapInterleaveLayout, "tensormap_interleave_layout"> {
  let assemblyFormat = "`<` $value `>`";
}

// Tensormap Swizzle Mode Enum
def TensormapSwizzleModeNoSwizzling : I32EnumCase<"NO_SWIZZLING", 0, "no_swizzling">;
def TensormapSwizzleMode32B : I32EnumCase<"B32", 1, "b32">;
def TensormapSwizzleMode64B : I32EnumCase<"B64", 2, "b64">;
def TensormapSwizzleMode128B : I32EnumCase<"B128", 3, "b128">;
def TensormapSwizzleMode96B : I32EnumCase<"B96", 4, "b96">;

def TensormapSwizzleMode : I32Enum<"TensormapSwizzleMode", "NVVM Tensormap Swizzle Mode",
  [TensormapSwizzleModeNoSwizzling, 
   TensormapSwizzleMode32B,
   TensormapSwizzleMode64B, 
   TensormapSwizzleMode128B,
   TensormapSwizzleMode96B]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapSwizzleModeAttr : EnumAttr<NVVM_Dialect,
  TensormapSwizzleMode, "tensormap_swizzle_mode"> {
  let assemblyFormat = "`<` $value `>`";
}

// Tensormap Swizzle Atomicity Enum
def TensormapSwizzleAtomicityB16 : I32EnumCase<"B16", 0, "b16">;
def TensormapSwizzleAtomicityB32 : I32EnumCase<"B32", 1, "b32">;
def TensormapSwizzleAtomicityB32FlipB8 : I32EnumCase<"B32_FLIP_8B", 2, "b32_flip_b8">;
def TensormapSwizzleAtomicityB64 : I32EnumCase<"B64", 3, "b64">;

def TensormapSwizzleAtomicity : 
  I32Enum<"TensormapSwizzleAtomicity", "NVVM Tensormap Swizzle Atomicity",
  [TensormapSwizzleAtomicityB16, TensormapSwizzleAtomicityB32,
   TensormapSwizzleAtomicityB32FlipB8, 
   TensormapSwizzleAtomicityB64]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapSwizzleAtomicityAttr : EnumAttr<NVVM_Dialect,
  TensormapSwizzleAtomicity, "tensormap_swizzle_atomicity"> {
  let assemblyFormat = "`<` $value `>`";
}

// Tensormap Fill Mode Enum
def TensormapFillModeZeroFill : I32EnumCase<"ZERO", 0, "zero">;
def TensormapFillModeOOBNaNFill : I32EnumCase<"OOB_NAN", 1, "oob_nan">;

def TensormapFillMode : I32Enum<"TensormapFillMode", "NVVM Tensormap Fill Mode",
  [TensormapFillModeZeroFill, TensormapFillModeOOBNaNFill]> {
    let cppNamespace = "::mlir::NVVM";
}

def TensormapFillModeAttr : EnumAttr<NVVM_Dialect,
  TensormapFillMode, "tensormap_fill_mode"> {
  let assemblyFormat = "`<` $value `>`";
}

def TensormapFieldValueAttr :
  AnyAttrOf<[TensormapElemtypeAttr, TensormapInterleaveLayoutAttr, 
    TensormapSwizzleModeAttr, TensormapSwizzleAtomicityAttr, 
    TensormapFillModeAttr]>;

def NVVM_TensormapReplaceOp : NVVM_VoidIntrinsicOp<"tensormap.replace",
    [NVVMRequiresSMaOrSMf<[90, 103], [100, 101, 110, 120]>]> {
  let summary = "Modifies a field of the tensor-map object";
  let description = [{
    The `nvvm.tensormap.replace` replaces the specified field of the tensor-map 
    object at the location specified by `addr` with a new value (specified by 
    `new_value` or `new_value_attr`).

    The `field` argument specifies the field of the tensor-map object to 
    replace.

    `new_value` is an `i32`/`i64` argument that specifies the new value to 
    replace the `field` with for the `global_address`, `rank`, `box_dim`, 
    `global_dim`, `global_stride`, and `element_stride` fields. It must be an 
    `i64` for the `global_address` and `global_stride` fields and `i32` for the 
    remaining fields.
    
    For `rank`, `new_value` must be one less than the desired tensor rank as 
    this field uses zero-based numbering.

    `new_value_attr` is an attribute that specifies the new value to replace 
    the `field` with for the `elemtype`, `interleave_layout`, `swizzle_mode`, 
    `swizzle_atomicity`, and `fill_mode` fields. It takes the place of 
    `new_value` for these fields. It must be a valid attribute corresponding to 
    the `field` type.

    The ordinal `ord` is an immediate integer argument that specifies the 
    ordinal of the `field` across the tensor which needs to be replaced and is 
    required only for the `box_dim`, `global_dim`, `global_stride`, and 
    `element_stride` fields.

    [For more information, see PTX ISA.](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-tensormap-replace)
  }];

  let hasVerifier = 1;
  let arguments = (ins
    TensormapFieldAttr:$field,
    AnyTypeOf<[LLVM_PointerGlobal, LLVM_PointerShared]>:$addr,
    Optional<AnyTypeOf<[I64, I32]>>:$new_value,
    OptionalAttr<ConfinedAttr<I32Attr, [IntMinValue<0>, IntMaxValue<4>]>>:$ord,
    OptionalAttr<TensormapFieldValueAttr>:$new_value_attr
  );

  let assemblyFormat = [{
    `field` `=` $field (`[` $ord^ `]`)? `,` `new_value` `=` ($new_value_attr^):($new_value)? `in` $addr attr-dict `:` type(operands)
  }];
}

//===----------------------------------------------------------------------===//
// NVVM target attribute.
//===----------------------------------------------------------------------===//

def NVVM_TargetAttr : NVVM_Attr<"NVVMTarget", "target", 
  [DeclareAttrInterfaceMethods<GPUTargetAttrVerifyInterface>]> {
  let description = [{
    GPU target attribute for controlling compilation of NVIDIA targets. All
    parameters decay into default values if not present.

    Examples:

    1. Target with default values.
    ```
      gpu.module @mymodule [#nvvm.target] attributes {...} {
        ...
      }
    ```

    2. Target with `sm_90` chip and fast math.
    ```
      gpu.module @mymodule [#nvvm.target<chip = "sm_90", flags = {fast}>] {
        ...
      }
    ```
  }];
  let parameters = (ins
    DefaultValuedParameter<"int", "2", "Optimization level to apply.">:$O,
    StringRefParameter<"Target triple.", "\"nvptx64-nvidia-cuda\"">:$triple,
    StringRefParameter<"Target chip.", "\"sm_75\"">:$chip,
    StringRefParameter<"Target chip features.", "\"\"">:$features,
    OptionalParameter<"DictionaryAttr", "Target specific flags.">:$flags,
    OptionalParameter<"ArrayAttr", "Files to link to the LLVM module.">:$link,
    DefaultValuedParameter<"bool", "true", "Perform SM version check on Ops.">:$verifyTarget
  );
  let assemblyFormat = [{
    (`<` struct($O, $triple, $chip, $features, $flags, $link, $verifyTarget)^ `>`)?
  }];
  let builders = [
    AttrBuilder<(ins CArg<"int", "2">:$optLevel,
                     CArg<"StringRef", "\"nvptx64-nvidia-cuda\"">:$triple,
                     CArg<"StringRef", "\"sm_75\"">:$chip,
                     CArg<"StringRef", "\"\"">:$features,
                     CArg<"DictionaryAttr", "nullptr">:$targetFlags,
                     CArg<"ArrayAttr", "nullptr">:$linkFiles,
                     CArg<"bool", "true">:$verifyTarget), [{
      return $_get($_ctxt, optLevel, triple, chip, features, targetFlags, linkFiles, verifyTarget);
    }]>];
  let skipDefaultBuilders = 1;
  let genVerifyDecl = 1;
  let extraClassDeclaration = [{
    bool hasFlag(StringRef flag) const;
    bool hasFastMath() const;
    bool hasFtz() const;
    bool hasCmdOptions() const;
    std::optional<mlir::NamedAttribute> getCmdOptions() const;
    LogicalResult verifyTarget(Operation *gpuModule);
  }];
  let extraClassDefinition = [{
    bool $cppClass::hasFlag(StringRef flag) const {
      if (DictionaryAttr flags = getFlags())
        return flags.get(flag) != nullptr;
      return false;
    }
    bool $cppClass::hasFastMath() const {
      return hasFlag("fast");
    }
    bool $cppClass::hasFtz() const {
      return hasFlag("ftz");
    }
    bool $cppClass::hasCmdOptions() const {
      return hasFlag("ptxas-cmd-options");
    }
    std::optional<mlir::NamedAttribute> $cppClass::getCmdOptions() const {
      return getFlags().getNamed("ptxas-cmd-options");
    }
  }];
}

#endif // NVVMIR_OPS