1 /*
2 * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25 package hat.test;
26
27 import hat.Accelerator;
28 import hat.ComputeContext;
29 import hat.KernelContext;
30 import static hat.KernelContext.*;
31 import hat.NDRange.Tile2D;
32 import hat.backend.Backend;
33 import hat.buffer.F16Array;
34 import hat.buffer.F32Array;
35 import hat.buffer.F32ArrayPadded;
36 import hat.test.annotation.HatTest;
37 import hat.test.exceptions.HATAssertionError;
38 import hat.test.exceptions.HATAsserts;
39 import hat.test.exceptions.HATExpectedPrecisionError;
40 import hat.types.F16;
41 import hat.types.Tensor;
42 import jdk.incubator.code.Reflect;
43
44 import java.lang.invoke.MethodHandles;
45 import java.util.Random;
46
47 import static hat.NDRange.Global2D;
48 import static hat.NDRange.Local2D;
49 import static hat.NDRange.NDRange2D;
50 import static hat.NDRange.Warp2D;
51
52 /**
53 * Test tensor operations in HAT.
54 *
55 * <p>How to run?</p>
56 * <p>For the CUDA backend:
57 * <code>
58 * HAT=SHOW_CODE java -cp hat/job.jar hat.java test ffi-cuda hat.test.TestTensors
59 * </code>
60 * </p>
61 *
62 * <p>For the OpenCL backend:
63 * <code>
64 * HAT=SHOW_CODE java -cp hat/job.jar hat.java test ffi-opencl hat.test.TestTensors
65 * </code>
66 * </p>
67 *
68 */
69 public class TestTensors {
70
71 @Reflect
72 public static void mxmTensorsColumnMajor(KernelContext kc, F16Array matrixA, F16Array matrixB, F32Array matrixC, int size) {
73 final int SHAPE = 16;
74 final int WMMA_M = SHAPE;
75 final int WMMA_N = SHAPE;
76 final int WMMA_K = SHAPE;
77 int warpM = GIX() / WRS();
78 int warpN = GIY();
79
80 final int lda = 1024;
81 final int ldb = 1024;
82 final int ldc = 1024;
83
84 var shape = Tensor.shape(WMMA_M, WMMA_N, WMMA_K);
85
86 // Initialize a tensor accumulator with zeros
87 Tensor acc = Tensor.zeros(shape, float.class);
88
89 for (int i = 0; i < size; i += WMMA_K) {
90 int aRow = warpM * WMMA_M;
91 int aCol = i;
92 int bRow = i;
93 int bCol = warpN * WMMA_N;
94
95 if (aRow < lda && aCol < lda && bRow < ldb && bCol < ldb) {
96 // Load data from matrix A with the specified shape using column-major into a tensor of FP16
97 Tensor tensorA = Tensor.loadF16(matrixA, aRow, aCol, lda, shape, Tensor.ofColumnMajor());
98
99 // Load data from matrix B with the specified shape using column-major into a tensor of FP16
100 Tensor tensorB = Tensor.loadF16(matrixB, bRow, bCol, ldb, shape, Tensor.ofColumnMajor());
101
102 // Perform the MMA operation:
103 // acc = tensorA * tensorB + acc
104 acc = Tensor.mma(tensorA, tensorB, acc);
105 }
106 }
107 int cRow = warpM * WMMA_M;
108 int cCol = warpN * WMMA_N;
109
110 // Store the resulting tensor into main memory using column-major layout.
111 if (cRow < size && cCol < size) {
112 // We operate with square matrices
113 Tensor.store(matrixC, cRow, cCol, acc, ldc, Tensor.ofColumnMajor());
114 }
115 }
116
117 @Reflect
118 public static void mxmTensorsColumnMajor(ComputeContext cc, F16Array matrixA, F16Array matrixB, F32Array matrixC, int globalSize) {
119 // The total number of threads is calculated as follows:
120 // [ (size / tile), (size / tile) ]
121 // If warpSize > 1, then each dimension using warp operations is multiplied by the value of the warp-size. This is architecture dependent, but the
122 // HAT runtime and HAT JIT compiler handle this automatically.
123
124 var ndRange = NDRange2D.of(Global2D.of(globalSize, globalSize),
125 Local2D.of(128, 4),
126 Tile2D.of(16, 16),
127 Warp2D.of(true, false));
128
129 cc.dispatchKernel(ndRange, kc -> mxmTensorsColumnMajor(kc, matrixA, matrixB, matrixC, globalSize));
130 }
131
132 @Reflect
133 public static void mxmTensorsRowColumnMajor(KernelContext kc, F16Array matrixA, F16Array matrixB, F32Array matrixC, int size) {
134
135 final int WMMA_M = 16;
136 final int WMMA_N = 16;
137 final int WMMA_K = 16;
138 int warpM = GIX() / WRS();
139 int warpN = GIY();
140
141 final int lda = 1024;
142 final int ldb = 1024;
143 final int ldc = 1024;
144
145 // We keep explicit constant in this version to check shape with ConstantOp
146 Tensor acc = Tensor.create(Tensor.shape(16, 16, 16), float.class);
147
148 Tensor.fill(acc, 0.0f);
149
150 for (int i = 0; i < size; i += WMMA_K) {
151 int aRow = warpM * WMMA_M;
152 int aCol = i;
153 int bRow = i;
154 int bCol = warpN * WMMA_N;
155 if (aRow < lda && aCol < lda && bRow < ldb && bCol < ldb) {
156 Tensor tensorA = Tensor.loadF16(matrixA, aRow, aCol, lda, Tensor.shape(16, 16, 16), Tensor.ofRowMajor());
157 Tensor tensorB = Tensor.loadF16(matrixB, bRow, bCol, ldb, Tensor.shape(16, 16, 16),Tensor.ofColumnMajor());
158 acc = Tensor.mma(tensorA, tensorB, acc);
159 }
160 }
161 int cRow = warpM * WMMA_M;
162 int cCol = warpN * WMMA_N;
163 if (cRow < size && cCol < size) {
164 // We operate with square matrices
165 Tensor.store(matrixC, cRow, cCol, acc, ldc, Tensor.ofColumnMajor());
166 }
167 }
168
169 @Reflect
170 public static void mxmTensorsRowColumnMajor(ComputeContext cc, F16Array matrixA, F16Array matrixB, F32Array matrixC, int globalSize) {
171 // The total number of threads is calculated as follows:
172 // [ (size / tile), (size / tile) ]
173 // If warpSize > 1, then each dimension using warp operations is multiplied by the value of the warp-size. This is architecture dependent, but the
174 // HAT runtime and HAT JIT compiler handle this automatically.
175
176 var ndRange = NDRange2D.of(Global2D.of(globalSize, globalSize),
177 Local2D.of(128, 4),
178 Tile2D.of(16, 16),
179 Warp2D.of(true, false));
180
181 cc.dispatchKernel(ndRange, kc -> mxmTensorsRowColumnMajor(kc, matrixA, matrixB, matrixC, globalSize));
182 }
183
184 @Reflect
185 public static void mxmTensorsRowMajor(KernelContext kc, F16Array matrixA, F16Array matrixB, F32ArrayPadded matrixC, int size) {
186 final int WMMA_M = 16;
187 final int WMMA_N = 16;
188 final int WMMA_K = 16;
189 int warpM = GIX() / WRS();
190 int warpN = GIY();
191
192 final int lda = 1024;
193 final int ldb = 1024;
194 final int ldc = 1024;
195
196 Tensor acc = Tensor.create(Tensor.shape(16, 16, 16), float.class);
197
198 Tensor.fill(acc, 0.0f);
199
200 for (int i = 0; i < size; i += WMMA_K) {
201 int aRow = warpM * WMMA_M;
202 int aCol = i;
203
204 int bRow = i;
205 int bCol = warpN * WMMA_N;
206
207 if (aRow < lda && aCol < lda && bRow < ldb && bCol < ldb) {
208
209 Tensor tensorA = Tensor.loadF16(matrixA, aRow, aCol, lda, Tensor.shape(16, 16, 16), Tensor.ofRowMajor());
210 Tensor tensorB = Tensor.loadF16(matrixB, bRow, bCol, ldb, Tensor.shape(16, 16, 16), Tensor.ofRowMajor());
211
212 // acc = tensorA * tensorB + acc
213 acc = Tensor.mma(tensorA, tensorB, acc);
214 }
215 }
216 int cRow = warpM * WMMA_M;
217 int cCol = warpN * WMMA_N;
218 if (cRow < size && cCol < size) {
219 Tensor.store(matrixC, cRow, cCol, acc, ldc, Tensor.ofRowMajor());
220 }
221 }
222
223 @Reflect
224 public static void mxmTensorsRowMajor( ComputeContext cc, F16Array matrixA, F16Array matrixB, F32ArrayPadded matrixC, int globalSize) {
225 var ndRange = NDRange2D.of(
226 Global2D.of(globalSize, globalSize),
227 Local2D.of(128, 4),
228 Tile2D.of(16, 16),
229 Warp2D.of(true, false));
230 cc.dispatchKernel(ndRange, kc -> mxmTensorsRowMajor(kc, matrixA, matrixB, matrixC, globalSize));
231 }
232
233 @Reflect
234 public static void mxmTensorsDefaultAccess( KernelContext kc, F16Array matrixA, F16Array matrixB, F32ArrayPadded matrixC, int size) {
235 final int sizeShape = 16;
236 final int WMMA_M = sizeShape;
237 final int WMMA_N = sizeShape;
238 final int WMMA_K = sizeShape;
239 int warpM = GIX() / WRS();
240 int warpN = GIY();
241
242 final int lda = 1024;
243 final int ldb = 1024;
244 final int ldc = 1024;
245
246 var shape = Tensor.shape(sizeShape, sizeShape, sizeShape);
247 Tensor acc = Tensor.zeros(shape, float.class);
248 for (int i = 0; i < size; i += WMMA_K) {
249 int aRow = warpM * WMMA_M;
250 int bCol = warpN * WMMA_N;
251 if (aRow < lda && i < lda && i < ldb && bCol < ldb) {
252 Tensor tensorA = Tensor.loadF16(matrixA, aRow, i, lda, shape);
253 Tensor tensorB = Tensor.loadF16(matrixB, i, bCol, ldb, shape);
254 acc = Tensor.mma(tensorA, tensorB, acc);
255 }
256 }
257 int cRow = warpM * WMMA_M;
258 int cCol = warpN * WMMA_N;
259 if (cRow < size && cCol < size) {
260 Tensor.store(matrixC, cRow, cCol, acc, ldc);
261 }
262 }
263
264 @Reflect
265 public static void mxmTensorsDefaultAccess( ComputeContext cc, F16Array matrixA, F16Array matrixB, F32ArrayPadded matrixC, int globalSize) {
266 var ndRange = NDRange2D.of(Global2D.of(globalSize, globalSize),
267 Local2D.of(128, 4),
268 Tile2D.of(16, 16),
269 Warp2D.of(true, false));
270 cc.dispatchKernel(ndRange, kc -> mxmTensorsDefaultAccess(kc, matrixA, matrixB, matrixC, globalSize));
271 }
272
273 private static void runSequentialColMajor(F16Array matrixA, F16Array matrixB, F32Array matrixC, final int size) {
274 for (int i = 0; i < size; i++) {
275 for (int j = 0; j < size; j++) {
276 float sum = 0.0f;
277 for (int k = 0; k < size; k++) {
278 F16 a = matrixA.array((long) k * size + i);
279 F16 b = matrixB.array((long) j * size + k);
280 F16 mul = F16.mul(a, b);
281 sum += F16.f16ToFloat(mul);
282 }
283 matrixC.array((long) j * size + i, sum);
284 }
285 }
286 }
287
288 private static void runSequentialRowAndColMajor(F16Array matrixA, F16Array matrixB, F32Array matrixC, final int size) {
289 for (int i = 0; i < size; i++) {
290 for (int j = 0; j < size; j++) {
291 float sum = 0.0f;
292 for (int k = 0; k < size; k++) {
293 F16 a = matrixA.array((long) i * size + k);
294 F16 b = matrixB.array((long) j * size + k);
295 F16 mul = F16.mul(a, b);
296 sum += F16.f16ToFloat(mul);
297 }
298 matrixC.array((long) j * size + i, sum);
299 }
300 }
301 }
302
303 private static void runSequentialRowMajor(F16Array matrixA, F16Array matrixB, F32Array matrixC, final int size) {
304 for (int i = 0; i < size; i++) {
305 for (int j = 0; j < size; j++) {
306 float sum = 0.0f;
307 for (int k = 0; k < size; k++) {
308 F16 a = matrixA.array((long) i * size + k);
309 F16 b = matrixB.array((long) k * size + j);
310 F16 mul = F16.mul(a, b);
311 sum += F16.f16ToFloat(mul);
312 }
313 matrixC.array((long) i * size + j, sum);
314 }
315 }
316 }
317
318 @HatTest
319 @Reflect
320 public void testTensor01() {
321 var accelerator = new Accelerator(MethodHandles.lookup(), Backend.FIRST);
322 final int size = 1024;
323
324 F16Array matrixAHalf = F16Array.create(accelerator, size * size);
325 F16Array matrixBHalf = F16Array.create(accelerator, size * size);
326 F32Array matrixC = F32Array.create(accelerator, size * size);
327 F32Array resultSequential = F32Array.create(accelerator, size * size);
328
329 Random r = new Random(19);
330 for (int j = 0; j < matrixAHalf.length(); j++) {
331 matrixAHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
332 matrixBHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
333 }
334
335 // Run multiple time
336 for (int i = 0; i < 10; i++) {
337 accelerator.compute(cc -> mxmTensorsColumnMajor(cc, matrixAHalf, matrixBHalf, matrixC, size));
338 }
339
340 runSequentialColMajor(matrixAHalf, matrixBHalf, resultSequential, size);
341
342 for (int i = 0; i < size; i++) {
343 for (int j = 0; j < size; j++) {
344 final int index = j * size + i;
345 float expectedValue = resultSequential.array(index);
346 float gotValue = matrixC.array(index);
347 try {
348 HATAsserts.assertEquals(expectedValue, gotValue, 0.1f);
349 } catch (HATAssertionError e) {
350 throw new HATExpectedPrecisionError("Expected: " + expectedValue + " but got " + gotValue);
351 }
352 }
353 }
354 }
355
356 @HatTest
357 @Reflect
358 public void testTensor02() {
359 var accelerator = new Accelerator(MethodHandles.lookup(), Backend.FIRST);
360 final int size = 1024;
361
362 F16Array matrixAHalf = F16Array.create(accelerator, size * size);
363 F16Array matrixBHalf = F16Array.create(accelerator, size * size);
364 F32Array matrixC = F32Array.create(accelerator, size * size);
365 F32Array resultSequential = F32Array.create(accelerator, size * size);
366
367 Random r = new Random(19);
368 for (int j = 0; j < matrixAHalf.length(); j++) {
369 matrixAHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
370 matrixBHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
371 }
372
373 accelerator.compute(cc -> mxmTensorsRowColumnMajor(cc, matrixAHalf, matrixBHalf, matrixC, size));
374
375 runSequentialRowAndColMajor(matrixAHalf, matrixBHalf, resultSequential, size);
376
377 for (int i = 0; i < size; i++) {
378 for (int j = 0; j < size; j++) {
379 final int index = j * size + i;
380 float expectedValue = resultSequential.array(index);
381 float gotValue = matrixC.array(index);
382 try {
383 HATAsserts.assertEquals(expectedValue, gotValue, 0.1f);
384 } catch (HATAssertionError e) {
385 throw new HATExpectedPrecisionError("Expected: " + expectedValue + " but got " + gotValue);
386 }
387 }
388 }
389 }
390
391 @HatTest
392 @Reflect
393 public void testTensor03() {
394
395 // To be able to run tensor-matmul in a row-major layout, we need to add padding.
396 // Thus, the result matrix must be of type F32ArrayPadded.
397
398 var accelerator = new Accelerator(MethodHandles.lookup(), Backend.FIRST);
399 final int size = 1024;
400
401 F16Array matrixAHalf = F16Array.create(accelerator, size * size);
402 F16Array matrixBHalf = F16Array.create(accelerator, size * size);
403 F32ArrayPadded matrixC = F32ArrayPadded.create(accelerator, size * size);
404 F32Array resultSequential = F32Array.create(accelerator, size * size);
405
406 Random r = new Random(19);
407 for (int j = 0; j < matrixAHalf.length(); j++) {
408 matrixAHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
409 matrixBHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
410 }
411 accelerator.compute(cc -> mxmTensorsRowMajor(cc, matrixAHalf, matrixBHalf, matrixC, size));
412 runSequentialRowMajor(matrixAHalf, matrixBHalf, resultSequential, size);
413
414 for (int i = 0; i < size; i++) {
415 for (int j = 0; j < size; j++) {
416 final int index = j * size + i;
417 float expectedValue = resultSequential.array(index);
418 float gotValue = matrixC.array(index);
419 try {
420 HATAsserts.assertEquals(expectedValue, gotValue, 0.1f);
421 } catch (HATAssertionError e) {
422 throw new HATExpectedPrecisionError("Expected: " + expectedValue + " but got " + gotValue);
423 }
424 }
425 }
426 }
427
428 @HatTest
429 @Reflect
430 public void testTensor04() {
431
432 // To be able to run tensor-matmul in a row-major layout, we need to add padding.
433 // Thus, the result matrix must be of type F32ArrayPadded.
434
435 var accelerator = new Accelerator(MethodHandles.lookup(), Backend.FIRST);
436 final int size = 1024;
437
438 F16Array matrixAHalf = F16Array.create(accelerator, size * size);
439 F16Array matrixBHalf = F16Array.create(accelerator, size * size);
440 F32ArrayPadded matrixC = F32ArrayPadded.create(accelerator, size * size);
441 F32Array resultSequential = F32Array.create(accelerator, size * size);
442
443 Random r = new Random(19);
444 for (int j = 0; j < matrixAHalf.length(); j++) {
445 matrixAHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
446 matrixBHalf.array(j).value(F16.floatToF16(r.nextFloat()).value());
447 }
448
449 accelerator.compute(cc -> mxmTensorsDefaultAccess(cc, matrixAHalf, matrixBHalf, matrixC, size));
450
451 runSequentialRowMajor(matrixAHalf, matrixBHalf, resultSequential, size);
452
453 for (int i = 0; i < size; i++) {
454 for (int j = 0; j < size; j++) {
455 final int index = j * size + i;
456 float expectedValue = resultSequential.array(index);
457 float gotValue = matrixC.array(index);
458 try {
459 HATAsserts.assertEquals(expectedValue, gotValue, 0.1f);
460 } catch (HATAssertionError e) {
461 throw new HATExpectedPrecisionError("Expected: " + expectedValue + " but got " + gotValue);
462 }
463 }
464 }
465 }
466 }