[mlir] [sparse] [gpu] adding transpose support to spmm spmv

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1816–1818

I think it'd be better to just inline these into the definition of GPU_MatTransposeMode (see SparseTensorSortKindEnum for an example using that style)

1826

It would be better to use GPU_Dialect.cppNamespace instead of this string. It's best to avoid using literals whenever possible, both for future-proofing (not that they're likely to change this one any time soon) and to help guide the reader (i.e., to make explicit when things are definitionally equal, rather than just happening to be contingently equal).

1864

I think it'd be better to name this modeA (or more ideally transposeModeA if there's any chance of there being other "modes" in the future). The term "op" is generally used for operations, which this isn't.

And ditto for the other ops.

1985

Since the GPU_MatTransposeModeAttr arguments seem to be attached to particular GPU_SparseHandle arguments, I think it'd be better to use a syntax that helps make that more clear. For example, you could use something like gpu.spmm_buffersize async [%dep] %env, %spmatA{%opA}, %dnmatB{%opB}, %dnmatC (where I chose to use curly braces since that's what's typically used for attributes).

And ditto for the other ops.

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1864

I felt hard to choose when I wrote this code. The reason of using op is that cusparse documentation name this argument as opA. Shall I still use modeA? Do you think it necessary to comment here that this is corresponding to the opA in cusparse documentation?

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1985

Also, you may want to make these attributes optional, where they default to NON_TRANSPOSE when not provided. Since doing so will help avoid needing to make most of the changes to Conversion/GPUCommon/GPUToLLVMConversion.cpp and the various mlir test files.

I think the cleanest way to do that would be to modify the definition of GPU_MatTransposeModeAttr to specify the default (so it can be shared by all the ops). Though you'll have to make sure to update the assemblyFormat of each op to allow dropping the curly braces when the mode isn't provided. And you'll probably need to add some additional builders for each op, so that they can be defaulted on the C++ side. (Ideally the standard builders should be able to handle this automatically, but I don't recall whether they do or not for this particular case.)

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1402–1403

To avoid repeating this idiom all over, you should define a helper function (cf., SparseTensor/Transforms/CodegenUtils for some examples of us doing this elsewhere). I'd say ditto for dType idiom, though those aren't touched by this patch; so feel free to leave this cleanup for a separate CL

1403

Style-guide says to use static_cast here, instead of the C-style cast.

1481

What does this todo actually mean? Cuz it looks like you are passing the transpose modes...

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1860

This will not be an SSA value, %opA, but an actual "literal" attribute, as in NON_TRANSPOSE, right?
So please update the examples accordingly

[K-Wu]

addressing wren's comments

[K-Wu]

rm TODOs already done

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1481

Good catch! I forgot to remove this after finishing this item

[K-Wu]

static_cast

[K-Wu]

reformat

[K-Wu]

upd examples in GPUOps.td

[K-Wu]

rename to TransposeOperator to address concern

[K-Wu]

renaming opA|B to tOpA|B

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1864

Okay let's see if this is better. I now use transposeOpA|B and tOpA|B.

[K-Wu]

fix compile error

[K-Wu]

add MatTransposeOp default value

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1985

This is somewhat addressed by the use of Default_GPU_MatTransposeOpAttr in the last commit

[K-Wu]

fix stale assemblyFormat in gpu.spmm

[K-Wu]

presumably closing two comments. Please reopen/make new comments if you would like to suggest otherwise.

[K-Wu]

rebase origin/main

[K-Wu]

fixing unresolved merge conflict

[K-Wu]

fixing unresolved merge conflict

[K-Wu]

fixing broken build

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1821

Note that these should always be kept in sync with

typedef enum {

CUSPARSE_OPERATION_NON_TRANSPOSE       = 0,
CUSPARSE_OPERATION_TRANSPOSE           = 1,
CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE = 2

} cusparseOperation_t;

I had a similar dilemma for the index types and the data types.

We definitely do not want to pull in cusparse.h values in here, but we also don't want to introduce inconsistencies later. I have no good solution for this yet, but at least document the requirement for keeping these consistent here

1854

add

NON_TRANSPOSE (the default), ...

since the doc page does not see the code above

1860

I like this syntax! Very neat!

Note that you get bonus points for a syntax where the default is never printed, only the non default cases
(okay for follow up,but we have similar policies for the sparse tensor type features)

1902

same on default

1951

default

2002

default

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

743

you can remove the transpose form TODO from L740 ;-) since you have this now

[K-Wu]

fixing test error

[K-Wu]

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

743

This diff introduces the transpose attribute but I haven't added rewrite rules to recognize these patterns in loop nest and set the matrix as transposed. Shall I keep this comment and address it in new diffs?

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1816

Even if you stick with the "op" name for the operation arguments, you should absolutely not use "op" here— because the "Op" suffix is explicitly and exclusively reserved for operations. Personally I'd go with TransposeMode since the "mat" part doesn't add any information/disambiguation. (If the "mat" prefix is there for consistency with cuSparse names, then I suppose it's okay)

1829–1830

You should instead add let defaultValue = "MatTransposeOp::NON_TRANSPOSE"; to the definition of GPU_MatTransposeOpAttr itself, since it's the natural default value for all uses of that EnumAttr. Whereas DefaultValuedAttr is more for situations where either (a) you don't have control over the original attribute definition and therefore can't change it, or (b) you have some op-specific default that's different from the usual default (or the usual lack thereof).

1864

I'm not a big fan of single-letter abbreviations like that. Single letters like "A"/"X" are okay because those are in some sense their full names, and the "spmat"/"dn" parts are adjectives applying to those names. (And, arguably, the "spmat"/"dn" parts can be dropped since that's already captured by the types themselves and so there would be no ambiguity when seeing getX or getA in the C++ code; though if that's what cuSparse calls them, then there's no harm in keeping the longer names.) Whereas "t" could mean anything, isn't part of a standardized naming scheme (like "A"/"X" are), and doesn't (imo) help to disambiguate what "op" means.

I totally get your motivation for going with "op", and I like where you're mind's at :) It's just unfortunate that "op" means something totally different in MLIR, so there's no good way to be consistent with both cuSparse and MLIR at the same time. Personally, I'd side with MLIR for this case, though I'd defer to our noble TL if he feels otherwise. My reasoning for resolving this case comes from a two-pronged argument about which side readers would be most familiar with, and hence is least likely to cause confusion. Since this is part of the MLIR ecosystem, folks must be at least somewhat familiar with MLIR. And since MLIR isn't ubiquitously popular, that (ironically) makes readers even more likely to be familiar with MLIR or trying to become so. (Whereas things which are ubiquitously popular often have dabblers who care more about The Other Thing.) Conversely, and I may well be wrong about this part, I don't think readers are as likely to be familiar with cuSparse, or rather as likely to have greater familiarity/comfortability with cuSparse than with MLIR.

Though now that you've mentioned that the name comes from cuSparse, I'm thinking I might know where they're coming from. Namely, I'm guessing they mean "op" in the sense of the category theory notation for duality ("A^{op}"), which means different things depending on what the "A" being dualized is: namely for real-matrices it's transpose, but for complex-matrices it's conjugate-transpose. But if that is indeed where they're coming from, then I feel like that's all the more reason to side with MLIR— since anyone who's familiar with the category theory notation will also be familiar with several other notations for these operations, since linear algebraists can't decide between the half-dozen different traditions for what to name them.

That said, if we do side with MLIR over cuSparse, I'm now thinking the better short-name would be "transA" rather than "modeA" (and if you're feeling more verbose, then the better middling-/long-names would be "transposeA" and "transposeModeA")

1885

You should make this whole chunk into an optional-group. I believe "( `{` $tOpA^ `}`)?" should be the correct syntax for that, though you may need to play around with it since sometimes the optional-group stuff doesn't parse in the way it seems like it should. For more information, cf: https://mlir.llvm.org/docs/DefiningDialects/AttributesAndTypes/#optional-and-default-valued-parameters

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1203

I think it'd read nicer to use the prefix "constant" rather than "genConstFrom".

And this function is a prime example of why the enum must not be named with "Op": since when reading this I got very confused about why the argument was already an operation rather than the enum/attr, even despite the fact that I'd just re-read through the code above!

1204

Why not just have this function use llvmInt32Type directly? (Since that's the type that's used everywhere, and since this function doesn't do any validation that the parameter is indeed a 32-bit integer type, and since no other type would really make much sense)

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

467

Do you actually need to pass this argument since it's the default value? (i.e., do the automatically generated builders allow skipping it) If so, then you should add a new builder definition which uses the default in lieu of requiring an argument here. (And ditto for all the other ops, of course)

mlir/test/Dialect/GPU/ops.mlir

338 ↗

(On Diff #525728)

If you adjust the assemblyFormat to make the transpose-mode optional like I showed, then you should be able to revert all these changes to the test files (since the generated printers ought to avoid printing the attribute whenever it's the default value)

[K-Wu]

attempting to remove Default_GPU_MatTransposeOpAttr

[K-Wu]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1864

That makes a lot of sense! I didn't know the cetegory theory notation and was confused when I saw cusparse uses op in their documentation. Thanks a lot for this very detailed reasoning. I will follow the MLIR naming convension as you mentioned then.

1885

Noted with great thanks! I really appreciate your mentioning the links and examples in all of your comments. These are great aid to me.

[K-Wu]

explaining default value

[K-Wu]

cleaning up int32type

[K-Wu]

fixing error and cleaning up style

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1816

nit, even though we will think about this, I would phrase it a bit different just in case this remain,
something like

To avoid coupling this dialect with cusparse.h specifics, we hardcoded magic literals in this enum.
Note that this should be kept in sync with cusparseOperation_t in cusparse.h:
typedef enum .....

// todo: find a proper way to keep them in sync?

[K-Wu]

fixing compile error

[K-Wu]

fixing compile errors and formatting

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1823

"todo" should be all-caps :)

As for the content of the todo itself: Assuming cusparse.h is part of the cuSparse library itself, at what point does that library become coupled with the mlir libraries? Is it passed at runtime like our sparse-tensor runtime-library, or is it linked in at some earlier stage? If it's linked in before runtime, then in the file that brings them together we can always add a static_assert to ensure they're consistent at that point. Whereas if it's dynamic-linked at runtime, then I'm not sure there's any good way of ensuring consistency that's worth the cost...

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

1202

If this works, that's cool; but if the callsites don't require knowing that this is the specific operation being returned, then it would probably be better to return Value (both to reduce coupling, and because there are a few places where the use of templates gets a bit squirrely about implicitly converting operations to their result values.)

1202–1204

I'm guessing this isn't going to pass the clang-format checks, so be sure to re-run git clang-format HEAD^

1203

This function doesn't need the specifics of ConversionPatternRewriter, so you should change this parameter to OpBuilder &builder instead.

1206

This doesn't look right to me. The getContext method already returns MLIRContext*, so you don't need the address-of operator there

1207

You should just use rewriter.getIntegerType(32) instead.

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

450

MLIR-style says variables should start with lowercase

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1816

please avoid breaking the 80-col. Just break, and continue on next line (okay to have the Note... more to the right

[K-Wu]

fixing compiling errors and addressing comments

[K-Wu]

try optional attr in builder

[K-Wu]

reverting

[aartbik]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1860

maybe use another TRANSPOSE mode in the example, since we won't print the default anymore

1908

same

1957

same, or perhaps even one as default, one different?

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

743

nah, I was just observing the overlap in TODOs now, but no biggie

[K-Wu]

attempting to add optional argument

[K-Wu]

attempting to do default value

[K-Wu]

attempting to make default value work

[K-Wu]

attempt

[wrengr]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1853–1855

I wonder if you could move this documentation to the definition of GPU_TransposeModeAttr itself, to avoid needing to repeat if for every op. (Just a nit, not a blocker)

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

467

Fwiw, I think this was better when explicitly passing modeA, rather than passing std::nullptr. But the thing I had in mind was to have something more like:

def GPU_SpMVBufferSizeOp : ... {
let builders = [
  OpBuilder<(ins
      Variadic<GPU_AsyncToken>:$asyncDependencies,
      GPU_SparseEnvHandle:$env,
      GPU_SparseSpMatHandle:$spmatA,
      GPU_SparseDnVecHandle:$dnX,
      GPU_SparseDnVecHandle:$dnY), [{
    auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
    return build($_builder, $_state, asyncDependencies, env, modeA, spmatA, dnX, dnY);
  }];
}

...or whatever massaging of that is necessary to get it to compile.

If you're having trouble getting this to work, just send me something on chat and we can try to work out the wrinkles :)

[K-Wu]

using the builder Wren suggested

[K-Wu]

fixing format and compile error in tablegen

[K-Wu]

fixing unsupported type in tablegen

[K-Wu]

fixed compile errors

[K-Wu]

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

467

I make it work! Let me know if this looks good to you.

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        Arg<GPU_SparseSpMatHandle>:$spmat);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) $spmat attr-dict
   }];
 }
 
+// To avoid coupling this dialect with cusparse.h specifics, we hardcoded magic 
+// literals in this enum. Note that this should be kept in sync with 
+// cusparseOperation_t in cusparse.h:
+// typedef enum {
+// CUSPARSE_OPERATION_NON_TRANSPOSE       = 0,
+// CUSPARSE_OPERATION_TRANSPOSE           = 1,
+// CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE = 2
+// } cusparseOperation_t;
+// TODO: find a proper way to keep them in sync?
+def GPU_TransposeMode : I32EnumAttr<"TransposeMode",
+    "transpose mode of sparse matrix supported by sparse tensor ops",
+    [
+      I32EnumAttrCase<"NON_TRANSPOSE", 0>, 
+      I32EnumAttrCase<"TRANSPOSE", 1>, 
+      I32EnumAttrCase<"CONJUGATE_TRANSPOSE", 2>,
+    ]> {
+      let genSpecializedAttr = 0;
+      let cppNamespace = GPU_Dialect.cppNamespace;
+}
+
+def GPU_TransposeModeAttr : EnumAttr<GPU_Dialect, GPU_TransposeMode,
+                                   "mat_transpose_mode">{
+  let defaultValue = "TransposeMode::NON_TRANSPOSE";
+}
+
 def GPU_SpMVBufferSizeOp : GPU_Op<"spmv_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SpMV operation";
   let description = [{
     The `gpu.spmv_buffer_size` operation returns the buffer size required
     to perform the SpMV operation on the given sparse matrix and dense vectors.
     The operation expects handles returned by previous sparse operations
     to construct an environment and the operands for SpMV.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token in addition to the environment.
 
+    The matrix arguments can also be associated with one of the following 
+    operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
+    is NON_TRANSPOSE.
+
     Example:
 
     ```mlir
-    %buffersz, %token = gpu.spmv_buffersize async [%dep] %env, %spmatA, %dnX, %dnY
+    %buffersz, %token = gpu.spmv_buffersize async [%dep] %env, %spmatA{NON_TRANSPOSE}, %dnX, %dnY
     ```
   }];
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
+                       GPU_TransposeModeAttr:$modeA,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnVecHandle:$dnX,
                        GPU_SparseDnVecHandle:$dnY);
   let results = (outs Res<Index>:$bufferSz, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
+  let builders = [OpBuilder<(ins
+      Variadic<GPU_AsyncToken>:$asyncDependencies,
+      GPU_SparseEnvHandle:$env,
+      GPU_SparseSpMatHandle:$spmatA,
+      GPU_SparseDnVecHandle:$dnX,
+      GPU_SparseDnVecHandle:$dnY), [{
+    auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
+    return build($_builder, $_state, asyncDependencies, env, modeA, spmatA, dnX,
+                 dnY);
+  }]];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA `,` $dnX `,` $dnY attr-dict
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY attr-dict
   }];
 }
 
 def GPU_SpMVOp : GPU_Op<"spmv", [GPU_AsyncOpInterface]> {
   let summary = "SpMV operation";
   let description = [{
     The `gpu.spmv` operation performs the SpMV operation on the given sparse matrix,
     dense vectors, and buffer.  The operation expects handles returned by previous
     sparse operations to construct an environment and the operands for SpMV. The
     buffer must have been allocated on the device.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token in addition to the environment.
 
+    The matrix arguments can also be associated with one of the following 
+    operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
+    is NON_TRANSPOSE.
+
     Example:
 
     ```mlir
-    %token = gpu.spmv async [%dep] %env, %spmatA, %dnX, %dnY : memref<?xf64>
+    %token = gpu.spmv async [%dep] %env, %spmatA{NON_TRANSPOSE}, %dnX, %dnY : memref<?xf64>
     ```
   }];
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
+                       GPU_TransposeModeAttr:$modeA,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnVecHandle:$dnX,
                        GPU_SparseDnVecHandle:$dnY,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let builders = [OpBuilder<(ins
+      Variadic<GPU_AsyncToken>:$asyncDependencies,
+      GPU_SparseEnvHandle:$env,
+      GPU_SparseSpMatHandle:$spmatA,
+      GPU_SparseDnVecHandle:$dnX,
+      GPU_SparseDnVecHandle:$dnY,
+      AnyMemRef:$buffer), [{
+    auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
+    return build($_builder, $_state, asyncDependencies, env, modeA, spmatA, dnX,
+                 dnY, buffer);
+  }]];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer)
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnX `,` $dnY `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 def GPU_SpMMBufferSizeOp : GPU_Op<"spmm_buffer_size", [GPU_AsyncOpInterface]> {
   let summary = "Precompute buffersize for SpMM operation";
   let description = [{
     The `gpu.spmm_buffer_size` operation returns the buffer size required
     to perform the SpMM operation on the given sparse and dense matrix.
     The operation expects handles returned by previous sparse operations
     to construct an environment and the operands for SpMM.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token in addition to the environment.
 
+    The matrix arguments can also be associated with one of the following 
+    operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
+    is NON_TRANSPOSE.
+
     Example:
 
     ```mlir
-    %buffersz, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA, %spmatB, %spmatC
+    %buffersz, %token = gpu.spmm_buffersize async [%dep] %env, %spmatA{NON_TRANSPOSE}, %dnmatB{NON_TRANSPOSE}, %dnmatC
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
+                       GPU_TransposeModeAttr:$modeA,
+                       GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnMatHandle:$dnmatB,
                        GPU_SparseDnMatHandle:$dnmatC);
   let results = (outs Res<Index>:$bufferSz, 
                       Optional<GPU_AsyncToken>:$asyncToken);
 
+  let builders = [OpBuilder<(ins
+      Variadic<GPU_AsyncToken>:$asyncDependencies,
+      GPU_SparseEnvHandle:$env,
+      GPU_SparseSpMatHandle:$spmatA,
+      GPU_SparseDnMatHandle:$dnmatB,
+      GPU_SparseDnMatHandle:$dnmatC), [{
+    auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
+    auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
+    return build($_builder, $_state, asyncDependencies, env, modeA, modeB, 
+                 spmatA, dnmatB, dnmatC);
+  }]];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA `,` $dnmatB `,` $dnmatC attr-dict
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC attr-dict
   }];
 }
 
 def GPU_SpMMOp : GPU_Op<"spmm", [GPU_AsyncOpInterface]> {
   let summary = "SpMM operation";
   let description = [{
     The `gpu.spmm` operation performs the SpMM operation on the given sparse and
     dense matrix, and buffer.  The operation expects handles returned by previous
     sparse operations to construct an environment and the operands for SpMM. The
     buffer must have been allocated on the device.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it returns a !gpu.async.token in addition to the environment.
 
+    The matrix arguments can also be associated with one of the following 
+    operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value
+    is NON_TRANSPOSE.
+
     Example:
 
     ```mlir
-    %token = gpu.spmm async [%dep] %env, %spmatA, %spmatB, %spmatC, %buffer
+    %token = gpu.spmm async [%dep] %env, %spmatA{NON_TRANSPOSE}, %dnmatB{NON_TRANSPOSE}, %dnmatC, %buffer
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                        GPU_SparseEnvHandle:$env,
+                       GPU_TransposeModeAttr:$modeA,
+                       GPU_TransposeModeAttr:$modeB,
                        GPU_SparseSpMatHandle:$spmatA,
                        GPU_SparseDnMatHandle:$dnmatB,
                        GPU_SparseDnMatHandle:$dnmatC,
                        AnyMemRef:$buffer);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
+  let builders = [OpBuilder<(ins
+      Variadic<GPU_AsyncToken>:$asyncDependencies,
+      GPU_SparseEnvHandle:$env,
+      GPU_SparseSpMatHandle:$spmatA,
+      GPU_SparseDnMatHandle:$dnmatB,
+      GPU_SparseDnMatHandle:$dnmatC,
+      AnyMemRef:$buffer), [{
+    auto modeA = gpu::TransposeMode::NON_TRANSPOSE;
+    auto modeB = gpu::TransposeMode::NON_TRANSPOSE;
+    return build($_builder, $_state, asyncDependencies, env, modeA, modeB, 
+                 spmatA, dnmatB, dnmatC, buffer);
+  }]];
+
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
-    $env `,` $spmatA `,` $dnmatB `,` $dnmatC `,` $buffer attr-dict `:` type($buffer)
+    $env `,` $spmatA (`{` $modeA^ `}`)? `,` $dnmatB (`{` $modeB^ `}`)? `,` $dnmatC `,` $buffer attr-dict `:` type($buffer)
   }];
 }
 
 #endif // GPU_OPS

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

FunctionCallBuilder createCsrCallBuilder = {
        llvmInt32Type, llvmPointerType /* void *stream */}};
   FunctionCallBuilder destroySpMatCallBuilder = {
       "mgpuDestroySpMat",
       llvmVoidType,
       {llvmPointerType, llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMVBufferSizeCallBuilder = {
       "mgpuSpMVBufferSize",
       llvmIntPtrType,
-      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
-       llvmInt32Type, llvmPointerType /* void *stream */}};
+      {llvmPointerType, llvmInt32Type, llvmPointerType, llvmPointerType,
+       llvmPointerType, llvmInt32Type, llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMVCallBuilder = {
       "mgpuSpMV",
       llvmVoidType,
-      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
-       llvmInt32Type, llvmPointerType, llvmPointerType /* void *stream */}};
+      {llvmPointerType, llvmInt32Type, llvmPointerType, llvmPointerType,
+       llvmPointerType, llvmInt32Type, llvmPointerType,
+       llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMMBufferSizeCallBuilder = {
       "mgpuSpMMBufferSize",
       llvmIntPtrType,
-      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
-       llvmInt32Type, llvmPointerType /* void *stream */}};
+      {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
+       llvmPointerType, llvmPointerType, llvmInt32Type,
+       llvmPointerType /* void *stream */}};
   FunctionCallBuilder spMMCallBuilder = {
       "mgpuSpMM",
       llvmVoidType,
-      {llvmPointerType, llvmPointerType, llvmPointerType, llvmPointerType,
-       llvmInt32Type, llvmPointerType, llvmPointerType /* void *stream */}};
+      {llvmPointerType, llvmInt32Type, llvmInt32Type, llvmPointerType,
+       llvmPointerType, llvmPointerType, llvmInt32Type, llvmPointerType,
+       llvmPointerType /* void *stream */}};
 };
 
 /// A rewrite pattern to convert gpu.host_register operations into a GPU runtime
 /// call. Currently it supports CUDA and ROCm (HIP).
 class ConvertHostRegisterOpToGpuRuntimeCallPattern
     : public ConvertOpToGpuRuntimeCallPattern<gpu::HostRegisterOp> {
 public:
   ConvertHostRegisterOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)

 static Type getSpMatElemType(Value spMat) {
   if (auto op = spMat.getDefiningOp<gpu::CreateCooOp>())
     return op.getValues().getType().cast<MemRefType>().getElementType();
   if (auto op = spMat.getDefiningOp<gpu::CreateCsrOp>())
     return op.getValues().getType().cast<MemRefType>().getElementType();
   llvm_unreachable("cannot find spmat def");
 }
 
+static Value genConstFrom(OpBuilder &builder, Location loc,
+                          gpu::TransposeMode mode) {
+  Type llvmInt32Type = builder.getIntegerType(32);
+  return builder.create<LLVM::ConstantOp>(loc, llvmInt32Type,
+                                          static_cast<int32_t>(mode));
+}
+
 LogicalResult ConvertCreateSparseEnvOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::CreateSparseEnvOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   auto stream = adaptor.getAsyncDependencies().front();

 
 LogicalResult ConvertSpMVBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMVBufferSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
+  auto modeA = genConstFrom(rewriter, loc, op.getModeA());
   Type dType = getSpMatElemType(op.getSpmatA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
   auto stream = adaptor.getAsyncDependencies().front();
   auto bufferSize =
       spMVBufferSizeCallBuilder
           .create(loc, rewriter,
-                  {adaptor.getEnv(), adaptor.getSpmatA(), adaptor.getDnX(),
-                   adaptor.getDnY(), dw, stream})
+                  {adaptor.getEnv(), modeA, adaptor.getSpmatA(),
+                   adaptor.getDnX(), adaptor.getDnY(), dw, stream})
           .getResult();
   rewriter.replaceOp(op, {bufferSize, stream});
   return success();
 }
 
 LogicalResult ConvertSpMVOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMVOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   Type dType = getSpMatElemType(op.getSpmatA());
+  auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
   auto stream = adaptor.getAsyncDependencies().front();
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
   spMVCallBuilder.create(loc, rewriter,
-                         {adaptor.getEnv(), adaptor.getSpmatA(),
+                         {adaptor.getEnv(), modeA, adaptor.getSpmatA(),
                           adaptor.getDnX(), adaptor.getDnY(), dw, pBuf,
                           stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 LogicalResult ConvertSpMMBufferSizeOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMMBufferSizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   Type dType = getSpMatElemType(op.getSpmatA());
+  auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
+  auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
   auto stream = adaptor.getAsyncDependencies().front();
   auto bufferSize =
       spMMBufferSizeCallBuilder
           .create(loc, rewriter,
-                  {adaptor.getEnv(), adaptor.getSpmatA(), adaptor.getDnmatB(),
-                   adaptor.getDnmatC(), dw, stream})
+                  {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(),
+                   adaptor.getDnmatB(), adaptor.getDnmatC(), dw, stream})
           .getResult();
   rewriter.replaceOp(op, {bufferSize, stream});
   return success();
 }
 
 LogicalResult ConvertSpMMOpToGpuRuntimeCallPattern::matchAndRewrite(
     gpu::SpMMOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (failed(areAllLLVMTypes(op, adaptor.getOperands(), rewriter)) ||
       failed(isAsyncWithOneDependency(rewriter, op)))
     return failure();
   Location loc = op.getLoc();
   Type dType = getSpMatElemType(op.getSpmatA());
   auto dw = rewriter.create<LLVM::ConstantOp>(loc, llvmInt32Type,
                                               dType.getIntOrFloatBitWidth());
+  auto modeA = genConstFrom(rewriter, loc, adaptor.getModeA());
+  auto modeB = genConstFrom(rewriter, loc, adaptor.getModeB());
   auto stream = adaptor.getAsyncDependencies().front();
   Value pBuf =
       MemRefDescriptor(adaptor.getBuffer()).allocatedPtr(rewriter, loc);
   if (!getTypeConverter()->useOpaquePointers())
     pBuf = rewriter.create<LLVM::BitcastOp>(loc, llvmPointerType, pBuf);
   spMMCallBuilder.create(loc, rewriter,
-                         {adaptor.getEnv(), adaptor.getSpmatA(),
+                         {adaptor.getEnv(), modeA, modeB, adaptor.getSpmatA(),
                           adaptor.getDnmatB(), adaptor.getDnmatC(), dw, pBuf,
                           stream});
   rewriter.replaceOp(op, {stream});
   return success();
 }
 
 template <typename T>
 static void addOpaquePointerConversion(LLVMTypeConverter &converter) {

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   Type envHandleTp = rewriter.getType<gpu::SparseEnvHandleType>();
   Type dnVecHandleTp = rewriter.getType<gpu::SparseDnVecHandleType>();
   Type spmatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
   Value token = genFirstWait(rewriter, loc);
   auto env =
       rewriter.create<gpu::CreateSparseEnvOp>(loc, envHandleTp, tokenTp, token);
   Value handle = env.getResult(0);
   token = env.getAsyncToken();
   Operation *spGenA =
       genSpMat(rewriter, loc, spmatHandleTp, tokenTp, token, szY, szX, nseA,
                rowA, colA, valA, isCOO, enableRT);
   Value spMatA = spGenA->getResult(0);
   token = spGenA->getResult(1);
   auto dvecX = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
                                                    token, vecX, szX);
   Value dnX = dvecX.getResult(0);
   token = dvecX.getAsyncToken();
   auto dvecY = rewriter.create<gpu::CreateDnVecOp>(loc, dnVecHandleTp, tokenTp,
                                                    token, vecY, szY);
   Value dnY = dvecY.getResult(0);
   token = dvecY.getAsyncToken();
 
   // Precompute buffersize for SpMV.
   auto bufferComp = rewriter.create<gpu::SpMVBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnX, dnY);
   Value bufferSz = bufferComp.getResult(0);
   token = bufferComp.getAsyncToken();
   auto buf = genAllocBuffer(rewriter, loc, bufferSz, token);
   Value buffer = buf.getResult(0);
   token = buf.getAsyncToken();
 
   // Perform the SpMV.
   auto spmvComp = rewriter.create<gpu::SpMVOp>(loc, tokenTp, token, handle,
auto bufferComp = rewriter.create<gpu::SpMMBufferSizeOp>(
       loc, indexTp, tokenTp, token, handle, spMatA, dnB, dnC);
   Value bufferSz = bufferComp.getResult(0);
   token = bufferComp.getAsyncToken();
   auto buf = genAllocBuffer(rewriter, loc, bufferSz, token);
   Value buffer = buf.getResult(0);
   token = buf.getAsyncToken();
 
   // Perform the SpMM.
-  auto spmmComp = rewriter.create<gpu::SpMMOp>(loc, tokenTp, token, handle,
-                                               spMatA, dnB, dnC, buffer);
+  auto spmmComp = rewriter.create<gpu::SpMMOp>(
+      loc, tokenTp, token, handle, spMatA, dnB, dnC, buffer);
   token = spmmComp.getAsyncToken();
 
   // Copy data back to host and free all the resoures.
   token = rewriter.create<gpu::DestroySpMatOp>(loc, tokenTp, token, spMatA)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnMatOp>(loc, tokenTp, token, dnB)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnMatOp>(loc, tokenTp, token, dnC)
LogicalResult matchAndRewrite(linalg::GenericOp op,
     bindDims(getContext(), i, j, k);
 
     // TODO: more robust patterns, tranposed versions, more kernels...
 
     // Recognize a SpMV kernel.
     if (numLoops == 2 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[0]) &&
         linalg::isReductionIterator(iteratorTypes[1]) &&
+        // TODO: add transposed {i, j}
         maps == infer({{i, j}, {j}, {i}}) && matchSumOfMultOfArgs(op)) {
       return rewriteSpMV(rewriter, op, enableRT);
     }
 
     // Recognize a SpMM kernel.
     if (numLoops == 3 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[0]) &&
         linalg::isParallelIterator(iteratorTypes[1]) &&
         linalg::isReductionIterator(iteratorTypes[2]) &&
+        // TODO: add transposed {i, k}, {k, j}
+        // TODO: maybe add transposed {i, j} in future
         maps == infer({{i, k}, {k, j}, {i, j}}) && matchSumOfMultOfArgs(op)) {
       return rewriteSpMM(rewriter, op, enableRT);
     }
 
     return failure();
   }
 
 private:

mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp

 }
 
 extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
 mgpuDestroySpMat(void *m, CUstream /*stream*/) {
   cusparseSpMatDescr_t mat = reinterpret_cast<cusparseSpMatDescr_t>(m);
   CUSPARSE_REPORT_IF_ERROR(cusparseDestroySpMat(mat))
 }
 
-extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMVBufferSize(
-    void *h, void *a, void *x, void *y, int32_t dw, CUstream /*stream*/) {
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
+mgpuSpMVBufferSize(void *h, int32_t ma, void *a, void *x, void *y, int32_t dw,
+                   CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
+  cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x);
   cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y);
   cudaDataType_t dtp = dataTp(dw);
   ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
-  CUSPARSE_REPORT_IF_ERROR(cusparseSpMV_bufferSize(
-      handle, CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, vecX, betap, vecY,
+  CUSPARSE_REPORT_IF_ERROR(
+      cusparseSpMV_bufferSize(handle, modeA, &alpha, matA, vecX, &beta, vecY,
                               dtp, CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
-extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(void *h, void *a, void *x,
-                                                   void *y, int32_t dw,
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMV(void *h, int32_t ma, void *a,
+                                                   void *x, void *y, int32_t dw,
                                                    void *buf,
                                                    CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
+  cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnVecDescr_t vecX = reinterpret_cast<cusparseDnVecDescr_t>(x);
   cusparseDnVecDescr_t vecY = reinterpret_cast<cusparseDnVecDescr_t>(y);
   cudaDataType_t dtp = dataTp(dw);
   ALPHABETA(dw, alpha, beta)
-  CUSPARSE_REPORT_IF_ERROR(
-      cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, vecX,
-                   betap, vecY, dtp, CUSPARSE_SPMV_ALG_DEFAULT, buf))
+  CUSPARSE_REPORT_IF_ERROR(cusparseSpMV(handle, modeA, &alpha, matA, vecX,
+                                        &beta, vecY, dtp,
+                                        CUSPARSE_SPMV_ALG_DEFAULT, buf))
 }
 
-extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t mgpuSpMMBufferSize(
-    void *h, void *a, void *b, void *c, int32_t dw, CUstream /*stream*/) {
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT intptr_t
+mgpuSpMMBufferSize(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c,
+                   int32_t dw, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
+  cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
+  cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c);
   cudaDataType_t dtp = dataTp(dw);
   ALPHABETA(dw, alpha, beta)
   size_t bufferSize = 0;
   CUSPARSE_REPORT_IF_ERROR(cusparseSpMM_bufferSize(
-      handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
-      CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, matB, betap, matC, dtp,
+      handle, modeA, modeB, &alpha, matA, matB, &beta, matC, dtp,
       CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize))
   return bufferSize == 0 ? 1 : bufferSize; // avoid zero-alloc
 }
 
-extern "C" MLIR_CUDA_WRAPPERS_EXPORT void mgpuSpMM(void *h, void *a, void *b,
-                                                   void *c, int32_t dw,
-                                                   void *buf,
-                                                   CUstream /*stream*/) {
+extern "C" MLIR_CUDA_WRAPPERS_EXPORT void
+mgpuSpMM(void *h, int32_t ma, int32_t mb, void *a, void *b, void *c, int32_t dw,
+         void *buf, CUstream /*stream*/) {
   cusparseHandle_t handle = reinterpret_cast<cusparseHandle_t>(h);
+  cusparseOperation_t modeA = static_cast<cusparseOperation_t>(ma);
+  cusparseOperation_t modeB = static_cast<cusparseOperation_t>(mb);
   cusparseSpMatDescr_t matA = reinterpret_cast<cusparseSpMatDescr_t>(a);
   cusparseDnMatDescr_t matB = reinterpret_cast<cusparseDnMatDescr_t>(b);
   cusparseDnMatDescr_t matC = reinterpret_cast<cusparseDnMatDescr_t>(c);
   cudaDataType_t dtp = dataTp(dw);
   ALPHABETA(dw, alpha, beta)
-  CUSPARSE_REPORT_IF_ERROR(
-      cusparseSpMM(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
-                   CUSPARSE_OPERATION_NON_TRANSPOSE, alphap, matA, matB, betap,
-                   matC, dtp, CUSPARSE_SPMM_ALG_DEFAULT, buf))
+  CUSPARSE_REPORT_IF_ERROR(cusparseSpMM(handle, modeA, modeB, &alpha, matA,
+                                        matB, &beta, matC, dtp,
+                                        CUSPARSE_SPMM_ALG_DEFAULT, buf))
 }