[mlir][sparse][gpu] Impl 2:4 SpMM rewrite for linalg op w/ DENSE24 attr

[K-Wu]

add test case

[K-Wu]

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib_2to4.mlir

62

I am not familiar with the attribute. Shall we use a flag like attr_name =1/0, or a flag where attr_name means set and its non-presence means not set? Is the latter possible in MLIR?

[aartbik]

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib_2to4.mlir

62

You can simply assign a string attribute to any MLIR op. So that would be a possible way until we have 2:4 properly defined into our sparse tensor type

[K-Wu]

updating attr in the test case

[K-Wu]

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib_2to4.mlir

62

SG!

[aartbik]

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

252

can you move this to the "Library helper methods" where we have all the is/are tests on type
this section is really more for the genXXX methods

also, this should probably use the isAllDense() or hasEncoding() utils from SparseTensorType

678

remove comment altogether

1021

use LLVM preferred style

if (op->getAttr("DENSE24"))

rewrite2To4SpMM(rewriter, op);

rewriteSpMM(rewriter, op, enableRT);

[K-Wu]

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

252

Addressed. Please check if it works.

[K-Wu]

addressing comments

[aartbik]

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

40

Oh, I meant to the section near L378, where the other type testers are

and, can't this be simply

return (getSparseTensorType(type).isAllDense())

or

return (getSparseTensorType(type).hasEncoding()

[K-Wu]

addressing comments

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

 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
 /// Marks the given top module as a GPU container module.
 static void markAsGPUContainer(ModuleOp topModule) {
   topModule->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
                      UnitAttr::get(topModule->getContext()));
 }
 
 /// Constructs a new GPU module (for GPU kernels) inside the given top module,
 /// or returns an existing GPU module if one was built previously.
 static gpu::GPUModuleOp genGPUModule(OpBuilder &builder, ModuleOp topModule) {
if (i == base) {
           genCopyMemRef(builder, loc, buffers[0], args[i], kernelToken);
     } else {
       firstToken = genFirstWait(builder, loc);
     }
     tokens.push_back(genDeallocMemRef(builder, loc, args[i], firstToken));
   }
 }
 
 /// Constructs code for new GPU kernel.
 static void genGPUCode(PatternRewriter &rewriter, gpu::GPUFuncOp gpuFunc,
                        scf::ParallelOp forallOp,
                        SmallVectorImpl<Value> &constants,
                        SmallVectorImpl<Value> &scalars,
                        SmallVectorImpl<Value> &buffers) {
   Location loc = gpuFunc->getLoc();
   Block &block = gpuFunc.getBody().front();
   rewriter.setInsertionPointToStart(&block);
if (auto unOp = other.getDefiningOp<sparse_tensor::UnaryOp>()) {
           redOp.getRegion().front().getTerminator());
       return matchMulOfArgs(op.getBlock(), yieldUn.getOperand(0)) &&
              matchAddOfArgs(&redOp.getRegion().front(), yieldRed.getOperand(0));
     }
   }
   return false;
 }
 
+/// Determines if the given value is a dense tensor instead of a sparse one.
+static bool isDenseTensor(Value v) {
+  return (sparse_tensor::getSparseTensorType(v).isAllDense());
+}
+
 /// Test for sorted COO with suitable data and coordinates types.
 static bool isAdmissibleCOO(SparseTensorType &aTp) {
   return aTp.isCompressedLvl(0) && aTp.isOrderedLvl(0) && !aTp.isUniqueLvl(0) &&
          aTp.isSingletonLvl(1) && aTp.isOrderedLvl(1) && aTp.isUniqueLvl(1) &&
          (aTp.getElementType().isF64() || aTp.getElementType().isF32()) &&
          (aTp.getCrdWidth() == 0 || aTp.getCrdWidth() == 32 ||
           aTp.getCrdWidth() == 64);
 }
static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Done.
   rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
   return success();
 }
 
+// Match and rewrite 2:4 SpMM kernels.
+static LogicalResult rewrite2To4SpMM(PatternRewriter &rewriter,
+                                     linalg::GenericOp op) {
+  Location loc = op.getLoc();
+  Value A = op.getOperand(0);
+  Value B = op.getOperand(1);
+  Value C = op.getOperand(2); // we have C = AB
+  SmallVector<Value> tokens;
+
+  // All input should be dense tensors.
+  if (!isDenseTensor(A) || !isDenseTensor(B) || !isDenseTensor(C))
+    return failure();
+
+  Value bufA = genTensorToMemref(rewriter, loc, A);
+  Value matA = genAllocCopy(rewriter, loc, bufA, tokens);
+  Value bufB = genTensorToMemref(rewriter, loc, B);
+  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
+  Value bufC = genTensorToMemref(rewriter, loc, C);
+  Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
+  genBlockingWait(rewriter, loc, tokens);
+  tokens.clear();
+  Value szm = linalg::createOrFoldDimOp(rewriter, loc, matA, 0);
+  Value szk = linalg::createOrFoldDimOp(rewriter, loc, matB, 0);
+  Value szn = linalg::createOrFoldDimOp(rewriter, loc, matC, 1);
+
+  Type indexTp = rewriter.getIndexType();
+  Type dnTensorHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
+  Type spMatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
+  Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
+  Value token = genFirstWait(rewriter, loc);
+  Operation *spGenA = rewriter.create<gpu::Create2To4SpMatOp>(
+      loc, spMatHandleTp, tokenTp, token, szm, szk, matA);
+
+  Value spMatA = spGenA->getResult(0);
+  token = spGenA->getResult(1);
+  auto dmatB = rewriter.create<gpu::CreateDnTensorOp>(
+      loc, dnTensorHandleTp, tokenTp, token, matB,
+      SmallVector<Value>{szk, szn});
+  Value dnB = dmatB.getResult(0);
+  token = dmatB.getAsyncToken();
+  auto dmatC = rewriter.create<gpu::CreateDnTensorOp>(
+      loc, dnTensorHandleTp, tokenTp, token, matC,
+      SmallVector<Value>{szm, szn});
+  Value dnC = dmatC.getResult(0);
+  token = dmatC.getAsyncToken();
+
+  auto dmatCType = llvm::cast<ShapedType>(matC.getType()).getElementType();
+
+  // Precompute buffersize for SpMM.
+  SmallVector<Type> bufferTypes_{indexTp, indexTp, indexTp};
+  TypeRange bufferTypes(bufferTypes_);
+  auto bufferComp = rewriter.create<gpu::SpMMBufferSizeOp>(
+      loc, bufferTypes, tokenTp, token, gpu::TransposeMode::NON_TRANSPOSE,
+      gpu::TransposeMode::NON_TRANSPOSE, spMatA, dnB, dnC,
+      /*computeType=*/dmatCType);
+
+  token = bufferComp.getAsyncToken();
+  Value bufferSz = bufferComp.getResult(0);
+  auto buf = genAllocBuffer(rewriter, loc, bufferSz, token);
+  Value buffer = buf.getResult(0);
+  token = buf.getAsyncToken();
+
+  Value bufferSz2 = bufferComp.getResult(1);
+  auto buf2 = genAllocBuffer(rewriter, loc, bufferSz2, token);
+  Value buffer2 = buf2.getResult(0);
+  token = buf2.getAsyncToken();
+
+  Value bufferSz3 = bufferComp.getResult(2);
+  auto buf3 = genAllocBuffer(rewriter, loc, bufferSz3, token);
+  Value buffer3 = buf3.getResult(0);
+  token = buf3.getAsyncToken();
+
+  auto dnCType = llvm::cast<ShapedType>(matC.getType()).getElementType();
+
+  // Perform the SpMM.
+  auto spmmComp = rewriter.create<gpu::SpMMOp>(
+      loc, tokenTp, token, spMatA, dnB, dnC, /*computeType=*/dnCType,
+      SmallVector<Value>{buffer, buffer2, buffer3});
+  token = spmmComp.getAsyncToken();
+
+  // Copy data back to host and free all the resources.
+  token = rewriter.create<gpu::DestroySpMatOp>(loc, tokenTp, token, spMatA)
+              .getAsyncToken();
+  token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnB)
+              .getAsyncToken();
+  token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnC)
+              .getAsyncToken();
+  SmallVector<Value> newDynamicSizes;
+
+  token = genDeallocMemRef(rewriter, loc, buffer, token);
+  token = genDeallocMemRef(rewriter, loc, buffer2, token);
+  token = genDeallocMemRef(rewriter, loc, buffer3, token);
+  token = genDeallocMemRef(rewriter, loc, matA, token);
+  token = genDeallocMemRef(rewriter, loc, matB, token);
+  token = genCopyMemRef(rewriter, loc, bufC, matC, token);
+  token = genDeallocMemRef(rewriter, loc, matC, token);
+  tokens.push_back(token);
+  genBlockingWait(rewriter, loc, tokens);
+  tokens.clear();
+  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
+  return success();
+}
+
 /// Match and rewrite SDDMM kernel.
 static LogicalResult rewriteSDDMM(PatternRewriter &rewriter,
                                   linalg::GenericOp op, bool enableRT) {
   Location loc = op.getLoc();
   Value a = op.getOperand(0);
   Value b = op.getOperand(1);
   Value c = op.getOperand(2);
   SmallVector<Value> tokens;
LogicalResult matchAndRewrite(linalg::GenericOp op,
     // 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)) {
+      if (op->getAttr("DENSE24"))
+        return rewrite2To4SpMM(rewriter, op);
+
       return rewriteSpMM(rewriter, op, enableRT);
     }
 
     // Recognize a SDDMM 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

mlir/test/Dialect/SparseTensor/GPU/gpu_matmul_lib_2to4.mlir

+// RUN: mlir-opt %s --linalg-generalize-named-ops \
+// RUN:             --sparsification="enable-gpu-libgen" | FileCheck %s
+
+// CHECK-LABEL:   func.func @matmul(
+// CHECK-SAME:                      %[[VAL_0:.*0]]: tensor<?x?xf16>,
+// CHECK-SAME:                      %[[VAL_1:.*1]]: tensor<?x?xf16>,
+// CHECK-SAME:                      %[[VAL_2:.*2]]: tensor<?x?xf16>) -> tensor<?x?xf16> {
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_5:.*]] = bufferization.to_memref %[[VAL_0]] : memref<?x?xf16>
+// CHECK:           %[[VAL_6:.*]] = gpu.wait async
+// CHECK:           %[[VAL_7:.*]] = memref.dim %[[VAL_5]], %[[VAL_3]] : memref<?x?xf16>
+// CHECK:           %[[VAL_8:.*]] = memref.dim %[[VAL_5]], %[[VAL_4]] : memref<?x?xf16>
+// CHECK:           %[[VAL_9:.*]], %[[VAL_10:.*]] = gpu.alloc async {{\[}}%[[VAL_6]]] (%[[VAL_7]], %[[VAL_8]]) : memref<?x?xf16>
+// CHECK:           %[[VAL_11:.*]] = gpu.memcpy async {{\[}}%[[VAL_10]]] %[[VAL_9]], %[[VAL_5]] : memref<?x?xf16>, memref<?x?xf16>
+// CHECK:           %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?xf16>
+// CHECK:           %[[VAL_13:.*]] = gpu.wait async
+// CHECK:           %[[VAL_14:.*]] = memref.dim %[[VAL_12]], %[[VAL_3]] : memref<?x?xf16>
+// CHECK:           %[[VAL_15:.*]] = memref.dim %[[VAL_12]], %[[VAL_4]] : memref<?x?xf16>
+// CHECK:           %[[VAL_16:.*]], %[[VAL_17:.*]] = gpu.alloc async {{\[}}%[[VAL_13]]] (%[[VAL_14]], %[[VAL_15]]) : memref<?x?xf16>
+// CHECK:           %[[VAL_18:.*]] = gpu.memcpy async {{\[}}%[[VAL_17]]] %[[VAL_16]], %[[VAL_12]] : memref<?x?xf16>, memref<?x?xf16>
+// CHECK:           %[[VAL_19:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?x?xf16>
+// CHECK:           %[[VAL_20:.*]] = gpu.wait async
+// CHECK:           %[[VAL_21:.*]] = memref.dim %[[VAL_19]], %[[VAL_3]] : memref<?x?xf16>
+// CHECK:           %[[VAL_22:.*]] = memref.dim %[[VAL_19]], %[[VAL_4]] : memref<?x?xf16>
+// CHECK:           %[[VAL_23:.*]], %[[VAL_24:.*]] = gpu.alloc async {{\[}}%[[VAL_20]]] (%[[VAL_21]], %[[VAL_22]]) : memref<?x?xf16>
+// CHECK:           %[[VAL_25:.*]] = gpu.memcpy async {{\[}}%[[VAL_24]]] %[[VAL_23]], %[[VAL_19]] : memref<?x?xf16>, memref<?x?xf16>
+// CHECK:           gpu.wait {{\[}}%[[VAL_11]], %[[VAL_18]], %[[VAL_25]]]
+// CHECK:           %[[VAL_26:.*]] = memref.dim %[[VAL_9]], %[[VAL_3]] : memref<?x?xf16>
+// CHECK:           %[[VAL_27:.*]] = memref.dim %[[VAL_16]], %[[VAL_3]] : memref<?x?xf16>
+// CHECK:           %[[VAL_28:.*]] = memref.dim %[[VAL_23]], %[[VAL_4]] : memref<?x?xf16>
+// CHECK:           %[[VAL_29:.*]] = gpu.wait async
+// CHECK:           %[[VAL_30:.*]], %[[VAL_31:.*]] = gpu.create_2to4_spmat async {{\[}}%[[VAL_29]]] %[[VAL_26]], %[[VAL_27]], %[[VAL_9]] : memref<?x?xf16>
+// CHECK:           %[[VAL_32:.*]], %[[VAL_33:.*]] = gpu.create_dn_tensor async {{\[}}%[[VAL_31]]] %[[VAL_16]], %[[VAL_27]], %[[VAL_28]] : index, index into memref<?x?xf16>
+// CHECK:           %[[VAL_34:.*]], %[[VAL_35:.*]] = gpu.create_dn_tensor async {{\[}}%[[VAL_33]]] %[[VAL_23]], %[[VAL_26]], %[[VAL_28]] : index, index into memref<?x?xf16>
+// CHECK:           %[[VAL_36:.*]]:3, %[[VAL_37:.*]] = gpu.spmm_buffer_size async {{\[}}%[[VAL_35]]] %[[VAL_30]], %[[VAL_32]], %[[VAL_34]] : index, index, index into f16
+// CHECK:           %[[VAL_38:.*]], %[[VAL_39:.*]] = gpu.alloc async {{\[}}%[[VAL_37]]] (%[[VAL_36]]#0) : memref<?xi8>
+// CHECK:           %[[VAL_40:.*]], %[[VAL_41:.*]] = gpu.alloc async {{\[}}%[[VAL_39]]] (%[[VAL_36]]#1) : memref<?xi8>
+// CHECK:           %[[VAL_42:.*]], %[[VAL_43:.*]] = gpu.alloc async {{\[}}%[[VAL_41]]] (%[[VAL_36]]#2) : memref<?xi8>
+// CHECK:           %[[VAL_44:.*]] = gpu.spmm async {{\[}}%[[VAL_43]]] %[[VAL_30]], %[[VAL_32]], %[[VAL_34]], %[[VAL_38]], %[[VAL_40]], %[[VAL_42]] : memref<?xi8>, memref<?xi8>, memref<?xi8> into f16
+// CHECK:           %[[VAL_45:.*]] = gpu.destroy_sp_mat async {{\[}}%[[VAL_44]]] %[[VAL_30]]
+// CHECK:           %[[VAL_46:.*]] = gpu.destroy_dn_tensor async {{\[}}%[[VAL_45]]] %[[VAL_32]]
+// CHECK:           %[[VAL_47:.*]] = gpu.destroy_dn_tensor async {{\[}}%[[VAL_46]]] %[[VAL_34]]
+// CHECK:           %[[VAL_48:.*]] = gpu.dealloc async {{\[}}%[[VAL_47]]] %[[VAL_38]] : memref<?xi8>
+// CHECK:           %[[VAL_49:.*]] = gpu.dealloc async {{\[}}%[[VAL_48]]] %[[VAL_40]] : memref<?xi8>
+// CHECK:           %[[VAL_50:.*]] = gpu.dealloc async {{\[}}%[[VAL_49]]] %[[VAL_42]] : memref<?xi8>
+// CHECK:           %[[VAL_51:.*]] = gpu.dealloc async {{\[}}%[[VAL_50]]] %[[VAL_9]] : memref<?x?xf16>
+// CHECK:           %[[VAL_52:.*]] = gpu.dealloc async {{\[}}%[[VAL_51]]] %[[VAL_16]] : memref<?x?xf16>
+// CHECK:           %[[VAL_53:.*]] = gpu.memcpy async {{\[}}%[[VAL_52]]] %[[VAL_19]], %[[VAL_23]] : memref<?x?xf16>, memref<?x?xf16>
+// CHECK:           %[[VAL_54:.*]] = gpu.dealloc async {{\[}}%[[VAL_53]]] %[[VAL_23]] : memref<?x?xf16>
+// CHECK:           gpu.wait {{\[}}%[[VAL_54]]]
+// CHECK:           %[[VAL_55:.*]] = bufferization.to_tensor %[[VAL_19]] : memref<?x?xf16>
+// CHECK:           return %[[VAL_55]] : tensor<?x?xf16>
+// CHECK:         }
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+module {
+  func.func @matmul(%arg0: tensor<?x?xf16>, %arg1: tensor<?x?xf16>, %arg2: tensor<?x?xf16>) -> tensor<?x?xf16> {
+    %0 = linalg.generic { DENSE24, indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<?x?xf16>, tensor<?x?xf16>) outs(%arg2 : tensor<?x?xf16>) {
+    ^bb0(%in: f16, %in_0: f16, %out: f16):
+      %1 = arith.mulf %in, %in_0 : f16
+      %2 = arith.addf %out, %1 : f16
+      linalg.yield %2 : f16
+    } -> tensor<?x?xf16>
+    return %0 : tensor<?x?xf16>
+  }
+}
 No newline at end of file